Note: If undo_transform=False then im_h and im_w are allowed to be None.
(dets_out, img, h, w, undo_transform=True, class_color=False, mask_alpha=0.45, fps_str='')
| 133 | color_cache = defaultdict(lambda: {}) |
| 134 | |
| 135 | def prep_display(dets_out, img, h, w, undo_transform=True, class_color=False, mask_alpha=0.45, fps_str=''): |
| 136 | """ |
| 137 | Note: If undo_transform=False then im_h and im_w are allowed to be None. |
| 138 | """ |
| 139 | if undo_transform: |
| 140 | img_numpy = undo_image_transformation(img, w, h) |
| 141 | img_gpu = torch.Tensor(img_numpy).cuda() |
| 142 | else: |
| 143 | img_gpu = img / 255.0 |
| 144 | h, w, _ = img.shape |
| 145 | |
| 146 | with timer.env('Postprocess'): |
| 147 | save = cfg.rescore_bbox |
| 148 | cfg.rescore_bbox = True |
| 149 | t = postprocess(dets_out, w, h, visualize_lincomb = args.display_lincomb, |
| 150 | crop_masks = args.crop, |
| 151 | score_threshold = args.score_threshold) |
| 152 | cfg.rescore_bbox = save |
| 153 | |
| 154 | with timer.env('Copy'): |
| 155 | idx = t[1].argsort(0, descending=True)[:args.top_k] |
| 156 | |
| 157 | if cfg.eval_mask_branch: |
| 158 | # Masks are drawn on the GPU, so don't copy |
| 159 | masks = t[3][idx] |
| 160 | classes, scores, boxes = [x[idx].cpu().numpy() for x in t[:3]] |
| 161 | |
| 162 | num_dets_to_consider = min(args.top_k, classes.shape[0]) |
| 163 | for j in range(num_dets_to_consider): |
| 164 | if scores[j] < args.score_threshold: |
| 165 | num_dets_to_consider = j |
| 166 | break |
| 167 | |
| 168 | # Quick and dirty lambda for selecting the color for a particular index |
| 169 | # Also keeps track of a per-gpu color cache for maximum speed |
| 170 | def get_color(j, on_gpu=None): |
| 171 | global color_cache |
| 172 | color_idx = (classes[j] * 5 if class_color else j * 5) % len(COLORS) |
| 173 | |
| 174 | if on_gpu is not None and color_idx in color_cache[on_gpu]: |
| 175 | return color_cache[on_gpu][color_idx] |
| 176 | else: |
| 177 | color = COLORS[color_idx] |
| 178 | if not undo_transform: |
| 179 | # The image might come in as RGB or BRG, depending |
| 180 | color = (color[2], color[1], color[0]) |
| 181 | if on_gpu is not None: |
| 182 | color = torch.Tensor(color).to(on_gpu).float() / 255. |
| 183 | color_cache[on_gpu][color_idx] = color |
| 184 | return color |
| 185 | |
| 186 | # First, draw the masks on the GPU where we can do it really fast |
| 187 | # Beware: very fast but possibly unintelligible mask-drawing code ahead |
| 188 | # I wish I had access to OpenGL or Vulkan but alas, I guess Pytorch tensor operations will have to suffice |
| 189 | if args.display_masks and cfg.eval_mask_branch and num_dets_to_consider > 0: |
| 190 | # After this, mask is of size [num_dets, h, w, 1] |
| 191 | masks = masks[:num_dets_to_consider, :, :, None] |
| 192 |
no test coverage detected