Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,343 Bytes
a249588 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import os
import cv2
import numpy as np
from sam2.distinctipy import get_colors
from pycocotools import mask as Mask
def batch_visualize_masks(args, image, masks_rle, image_kpts, bboxes_xyxy, dt_bboxes, gt_masks_raw, bbox_ious, mask_ious, image_path=None, mask_out=False, alpha=1.0):
# Decode dt_masks_rle
dt_masks = []
for mask_rle in masks_rle:
mask = Mask.decode(mask_rle)
dt_masks.append(mask)
dt_masks = np.array(dt_masks)
# Decode gt_masks_raw
gt_masks = []
for gt_mask in gt_masks_raw:
if gt_mask is None:
gt_masks.append(np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8))
else:
# gt_mask_rle = Mask.frPyObjects(gt_mask, image.shape[0], image.shape[1])
# gt_mask_rle = Mask.merge(gt_mask_rle)
gt_mask_rle = gt_mask
mask = Mask.decode(gt_mask_rle)
gt_masks.append(mask)
gt_masks = np.array(gt_masks)
# Generate random color for each mask
if mask_out:
dt_mask_image = dt_masks.max(axis=0)
dt_mask_image = (~ dt_mask_image.astype(bool)).astype(np.uint8)
dt_mask_image = cv2.resize(dt_mask_image, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
dt_mask_image = image * dt_mask_image[:, :, None]
dt_mask_image = cv2.addWeighted(image, 1-alpha, dt_mask_image, alpha, 0)
else:
colors = (np.array(get_colors(dt_masks.shape[0])) * 255).astype(int)
# colors = np.random.randint(0, 255, (dt_masks.shape[0], 3))
# # Make sure no colors are too dark
# np.clip(colors, 50, 255, out=colors)
# Repeat masks to 3 channels
dt_masks = np.repeat(dt_masks[:, :, :, None], 3, axis=3)
gt_masks = np.repeat(gt_masks[:, :, :, None], 3, axis=3)
# Colorize masks
dt_masks = dt_masks * colors[:, None, None, :]
gt_masks = gt_masks * colors[:, None, None, :]
# # Remove masks that are too small
# dt_masks_area = dt_masks.any(axis=3).sum(axis=(1, 2))
# dt_masks[dt_masks_area < 300*300] = 0
# Collapse masks to 3 channels
dt_mask_image = dt_masks.max(axis=0)
gt_mask_image = gt_masks.max(axis=0)
# Convert to uint8
dt_mask_image = dt_mask_image.astype(np.uint8)
gt_mask_image = gt_mask_image.astype(np.uint8)
# Resize masks to image size
dt_mask_image = cv2.resize(dt_mask_image, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
gt_mask_image = cv2.resize(gt_mask_image, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
# Add masks to image
if not mask_out:
dt_mask_image = cv2.addWeighted(image, 0.6, dt_mask_image, 0.4, 0)
# Draw contours around the masks
for mask, color in zip(dt_masks, colors):
color = color.astype(int).tolist()
mask = mask.astype(np.uint8)
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(dt_mask_image, contours, -1, color, 1)
gt_mask_image = cv2.addWeighted(image, 0.6, gt_mask_image, 0.4, 0)
# Draw keypoints
if image_kpts is not None and not mask_out:
for instance_kpts, color in zip(image_kpts, colors):
color = tuple(color.astype(int).tolist())
for kpt in instance_kpts:
cv2.circle(dt_mask_image, kpt.astype(int)[:2], 3, color, -1)
cv2.circle(gt_mask_image, kpt.astype(int)[:2], 3, color, -1)
# Draw bboxes
if bboxes_xyxy is not None and not mask_out:
bboxes_xyxy = np.array(bboxes_xyxy)
dt_bboxes = np.array(dt_bboxes)
dt_bboxes[:, 2:] += dt_bboxes[:, :2]
for gt_bbox, dt_bbox, color, biou in zip(bboxes_xyxy, dt_bboxes, colors, bbox_ious):
color = tuple(color.astype(int).tolist())
gbox = gt_bbox.astype(int)
dbox = dt_bbox.astype(int)
cv2.rectangle(dt_mask_image, (dbox[0], dbox[1]), (dbox[2], dbox[3]), color, 2)
cv2.rectangle(gt_mask_image, (gbox[0], gbox[1]), (gbox[2], gbox[3]), color, 2)
# Write IOU on th etop-left corner of the bbox
# cv2.putText(dt_mask_image, "{:.2f}".format(biou), (dbox[0], dbox[1]-2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
# cv2.putText(gt_mask_image, "{:.2f}".format(biou), (gbox[0], gbox[1]-2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
# Save the image
bbox_ious = np.array(bbox_ious)
mask_ious = np.array(mask_ious)
if image_path is not None:
save_name = os.path.basename(image_path)
else:
save_name = "batch_bbox_{:06.2f}_mask_{:06.2f}_{:02d}kpts_{:06d}.jpg".format(
bbox_ious.mean(), mask_ious.mean(), args.num_pos_keypoints, np.random.randint(1000000),
)
if 'debug_folder' not in args:
args.debug_folder = "debug"
if mask_out:
cv2.imwrite(os.path.join(args.debug_folder, save_name), dt_mask_image)
else:
cv2.imwrite(os.path.join(args.debug_folder, save_name), np.hstack([gt_mask_image, dt_mask_image]))
|