|
import json |
|
import cv2 |
|
import numpy as np |
|
import os |
|
from torch.utils.data import Dataset |
|
from PIL import Image |
|
import cv2 |
|
from .data_utils import * |
|
from panopticapi.utils import rgb2id |
|
from PIL import Image |
|
from .base import BaseDataset |
|
|
|
class VIPSegDataset(BaseDataset): |
|
def __init__(self, image_dir, anno): |
|
self.image_root = image_dir |
|
self.anno_root = anno |
|
video_dirs = [] |
|
video_dirs = os.listdir(self.image_root) |
|
self.data = video_dirs |
|
self.size = (512,512) |
|
self.clip_size = (224,224) |
|
self.dynamic = 1 |
|
|
|
def __len__(self): |
|
return 30000 |
|
|
|
def check_region_size(self, image, yyxx, ratio, mode = 'max'): |
|
pass_flag = True |
|
H,W = image.shape[0], image.shape[1] |
|
H,W = H * ratio, W * ratio |
|
y1,y2,x1,x2 = yyxx |
|
h,w = y2-y1,x2-x1 |
|
if mode == 'max': |
|
if h > H or w > W: |
|
pass_flag = False |
|
elif mode == 'min': |
|
if h < H or w < W: |
|
pass_flag = False |
|
return pass_flag |
|
|
|
def get_sample(self, idx): |
|
video_name = self.data[idx] |
|
video_path = os.path.join(self.image_root, video_name) |
|
frames = os.listdir(video_path) |
|
|
|
|
|
min_interval = len(frames) // 100 |
|
start_frame_index = np.random.randint(low=0, high=len(frames) - min_interval) |
|
end_frame_index = start_frame_index + np.random.randint(min_interval, len(frames) - start_frame_index ) |
|
end_frame_index = min(end_frame_index, len(frames) - 1) |
|
|
|
|
|
ref_image_name = frames[start_frame_index] |
|
tar_image_name = frames[end_frame_index] |
|
ref_image_path = os.path.join(self.image_root, video_name, ref_image_name) |
|
tar_image_path = os.path.join(self.image_root, video_name, tar_image_name) |
|
|
|
ref_mask_path = ref_image_path.replace('images','panomasksRGB').replace('.jpg', '.png') |
|
tar_mask_path = tar_image_path.replace('images','panomasksRGB').replace('.jpg', '.png') |
|
|
|
|
|
ref_image = cv2.imread(ref_image_path) |
|
ref_image = cv2.cvtColor(ref_image, cv2.COLOR_BGR2RGB) |
|
|
|
tar_image = cv2.imread(tar_image_path) |
|
tar_image = cv2.cvtColor(tar_image, cv2.COLOR_BGR2RGB) |
|
|
|
ref_mask = np.array(Image.open(ref_mask_path).convert('RGB')) |
|
ref_mask = rgb2id(ref_mask) |
|
|
|
tar_mask = np.array(Image.open(tar_mask_path).convert('RGB')) |
|
tar_mask = rgb2id(tar_mask) |
|
|
|
ref_ids = np.unique(ref_mask) |
|
tar_ids = np.unique(tar_mask) |
|
|
|
common_ids = list(np.intersect1d(ref_ids, tar_ids)) |
|
common_ids = [ i for i in common_ids if i != 0 ] |
|
|
|
chosen_id = np.random.choice(common_ids) |
|
ref_mask = ref_mask == chosen_id |
|
tar_mask = tar_mask == chosen_id |
|
|
|
len_mask = len( self.check_connect( ref_mask.astype(np.uint8) ) ) |
|
assert len_mask == 1 |
|
|
|
item_with_collage = self.process_pairs(ref_image, ref_mask, tar_image, tar_mask) |
|
sampled_time_steps = self.sample_timestep() |
|
item_with_collage['time_steps'] = sampled_time_steps |
|
return item_with_collage |
|
|
|
def check_connect(self, mask): |
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) |
|
cnt_area = [cv2.contourArea(cnt) for cnt in contours] |
|
return cnt_area |
|
|
|
|
