| 69 | |
| 70 | |
| 71 | def demo_postprocess(outputs, img_size, p6=False): |
| 72 | |
| 73 | grids = [] |
| 74 | expanded_strides = [] |
| 75 | |
| 76 | if not p6: |
| 77 | strides = [8, 16, 32] |
| 78 | else: |
| 79 | strides = [8, 16, 32, 64] |
| 80 | |
| 81 | hsizes = [img_size[0] // stride for stride in strides] |
| 82 | wsizes = [img_size[1] // stride for stride in strides] |
| 83 | |
| 84 | for hsize, wsize, stride in zip(hsizes, wsizes, strides): |
| 85 | xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) |
| 86 | grid = np.stack((xv, yv), 2).reshape(1, -1, 2) |
| 87 | grids.append(grid) |
| 88 | shape = grid.shape[:2] |
| 89 | expanded_strides.append(np.full((*shape, 1), stride)) |
| 90 | |
| 91 | grids = np.concatenate(grids, 1) |
| 92 | expanded_strides = np.concatenate(expanded_strides, 1) |
| 93 | outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides |
| 94 | outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides |
| 95 | |
| 96 | return outputs |