Spaces:
Running
on
A10G
Running
on
A10G
File size: 3,385 Bytes
19a149b |
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 |
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 PIL import Image
from .base import BaseDataset
class MoseDataset(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 = 2
def __len__(self):
return 40000
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)
# Sampling frames
min_interval = len(frames) // 10
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)
# Get image path
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('JPEGImages','Annotations').replace('.jpg', '.png')
tar_mask_path = tar_image_path.replace('JPEGImages','Annotations').replace('.jpg', '.png')
# Read Image and Mask
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 = Image.open(ref_mask_path ).convert('P')
ref_mask= np.array(ref_mask)
tar_mask = Image.open(tar_mask_path ).convert('P')
tar_mask= np.array(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 ]
assert len(common_ids) > 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
|