|
|
import argparse |
|
|
import json |
|
|
import tqdm |
|
|
import cv2 |
|
|
import os |
|
|
import numpy as np |
|
|
from pycocotools import mask as mask_utils |
|
|
import random |
|
|
from PIL import Image |
|
|
from natsort import natsorted |
|
|
|
|
|
EVALMODE = "test" |
|
|
|
|
|
|
|
|
def blend_mask(input_img, binary_mask, alpha=0.5, color="g"): |
|
|
if input_img.ndim == 2: |
|
|
return input_img |
|
|
mask_image = np.zeros(input_img.shape, np.uint8) |
|
|
if color == "r": |
|
|
mask_image[:, :, 0] = 255 |
|
|
if color == "g": |
|
|
mask_image[:, :, 1] = 255 |
|
|
if color == "b": |
|
|
mask_image[:, :, 2] = 255 |
|
|
if color == "o": |
|
|
mask_image[:, :, 0] = 255 |
|
|
mask_image[:, :, 1] = 165 |
|
|
mask_image[:, :, 2] = 0 |
|
|
if color == "c": |
|
|
mask_image[:, :, 0] = 0 |
|
|
mask_image[:, :, 1] = 255 |
|
|
mask_image[:, :, 2] = 255 |
|
|
if color == "p": |
|
|
mask_image[:, :, 0] = 128 |
|
|
mask_image[:, :, 1] = 0 |
|
|
mask_image[:, :, 2] = 128 |
|
|
|
|
|
mask_image = mask_image * np.repeat(binary_mask[:, :, np.newaxis], 3, axis=2) |
|
|
blend_image = input_img[:, :, :].copy() |
|
|
pos_idx = binary_mask > 0 |
|
|
for ind in range(input_img.ndim): |
|
|
ch_img1 = input_img[:, :, ind] |
|
|
ch_img2 = mask_image[:, :, ind] |
|
|
ch_img3 = blend_image[:, :, ind] |
|
|
ch_img3[pos_idx] = alpha * ch_img1[pos_idx] + (1 - alpha) * ch_img2[pos_idx] |
|
|
blend_image[:, :, ind] = ch_img3 |
|
|
return blend_image |
|
|
|
|
|
|
|
|
def upsample_mask(mask, frame): |
|
|
H, W = frame.shape[:2] |
|
|
mH, mW = mask.shape[:2] |
|
|
|
|
|
if W > H: |
|
|
ratio = mW / W |
|
|
h = H * ratio |
|
|
diff = int((mH - h) // 2) |
|
|
if diff == 0: |
|
|
mask = mask |
|
|
else: |
|
|
mask = mask[diff:-diff] |
|
|
else: |
|
|
ratio = mH / H |
|
|
w = W * ratio |
|
|
diff = int((mW - w) // 2) |
|
|
if diff == 0: |
|
|
mask = mask |
|
|
else: |
|
|
mask = mask[:, diff:-diff] |
|
|
|
|
|
mask = cv2.resize(mask, (W, H)) |
|
|
return mask |
|
|
|
|
|
|
|
|
def downsample(mask, frame): |
|
|
H, W = frame.shape[:2] |
|
|
mH, mW = mask.shape[:2] |
|
|
|
|
|
mask = cv2.resize(mask, (W, H)) |
|
|
return mask |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
|
|
color = ['g', 'r', 'b', 'o', 'c', 'p'] |
|
|
data_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/handal_dataset_whisks/train" |
|
|
output_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/handal_dataset_whisks/vis" |
|
|
mask_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/handal_dataset_whisks/mask_predictions/train" |
|
|
|
|
|
|
|
|
takes_ids = os.listdir(data_path) |
|
|
|
|
|
for take_id in tqdm.tqdm(takes_ids): |
|
|
|
|
|
prediction_path = os.path.join(mask_path, take_id) |
|
|
if not os.path.exists(prediction_path): |
|
|
print(take_id) |
|
|
continue |
|
|
|
|
|
|
|
|
prediction_path = os.path.join(prediction_path, "rgb") |
|
|
|
|
|
file_names = natsorted(os.listdir(prediction_path)) |
|
|
idxs = [f.split(".")[0] for f in file_names] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
for id in idxs: |
|
|
frame_idx = id |
|
|
frame = cv2.imread( |
|
|
f"{data_path}/{take_id}/rgb/{frame_idx}.jpg" |
|
|
) |
|
|
mask = Image.open(f"{prediction_path}/{frame_idx}.png") |
|
|
mask = np.array(mask) |
|
|
unique_instances = np.unique(mask) |
|
|
unique_instances = unique_instances[unique_instances != 0] |
|
|
for i,instance_value in enumerate(unique_instances): |
|
|
binary_mask = (mask == instance_value).astype(np.uint8) |
|
|
binary_mask = cv2.resize(binary_mask, (frame.shape[1], frame.shape[0])) |
|
|
binary_mask = upsample_mask(binary_mask, frame) |
|
|
out = blend_mask(frame, binary_mask, color=color[0]) |
|
|
os.makedirs( |
|
|
f"{output_path}/{take_id}/obj_{i}", |
|
|
exist_ok=True, |
|
|
) |
|
|
cv2.imwrite( |
|
|
f"{output_path}/{take_id}/obj_{i}/{frame_idx}.jpg", |
|
|
out, |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
