(image)
| 297 | return masks_new |
| 298 | |
| 299 | def sam_encoder(image): |
| 300 | image = cv2.cvtColor(image[0].permute(1,2,0).numpy().astype(np.uint8), cv2.COLOR_BGR2RGB) |
| 301 | # pre-compute masks |
| 302 | masks_default, masks_s, masks_m, masks_l = mask_generator.generate(image) |
| 303 | # pre-compute postprocess |
| 304 | masks_default, masks_s, masks_m, masks_l = \ |
| 305 | masks_update(masks_default, masks_s, masks_m, masks_l, iou_thr=0.8, score_thr=0.7, inner_thr=0.5) |
| 306 | |
| 307 | def mask2segmap(masks, image): |
| 308 | seg_img_list = [] |
| 309 | seg_map = -np.ones(image.shape[:2], dtype=np.int32) |
| 310 | for i in range(len(masks)): |
| 311 | mask = masks[i] |
| 312 | seg_img = get_seg_img(mask, image) |
| 313 | pad_seg_img = cv2.resize(pad_img(seg_img), (224,224)) |
| 314 | seg_img_list.append(pad_seg_img) |
| 315 | |
| 316 | seg_map[masks[i]['segmentation']] = i |
| 317 | seg_imgs = np.stack(seg_img_list, axis=0) # b,H,W,3 |
| 318 | seg_imgs = (torch.from_numpy(seg_imgs.astype("float32")).permute(0,3,1,2) / 255.0).to('cuda') |
| 319 | |
| 320 | return seg_imgs, seg_map |
| 321 | |
| 322 | seg_images, seg_maps = {}, {} |
| 323 | seg_images['default'], seg_maps['default'] = mask2segmap(masks_default, image) |
| 324 | if len(masks_s) != 0: |
| 325 | seg_images['s'], seg_maps['s'] = mask2segmap(masks_s, image) |
| 326 | if len(masks_m) != 0: |
| 327 | seg_images['m'], seg_maps['m'] = mask2segmap(masks_m, image) |
| 328 | if len(masks_l) != 0: |
| 329 | seg_images['l'], seg_maps['l'] = mask2segmap(masks_l, image) |
| 330 | |
| 331 | # 0:default 1:s 2:m 3:l |
| 332 | return seg_images, seg_maps |
| 333 | |
| 334 | def seed_everything(seed_value): |
| 335 | random.seed(seed_value) |
no test coverage detected