Spaces:
Runtime error
Runtime error
File size: 2,523 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 |
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 .base import BaseDataset
import albumentations as A
class FashionTryonDataset(BaseDataset):
def __init__(self, image_dir):
self.image_root = image_dir
self.data =os.listdir(self.image_root)
self.size = (512,512)
self.clip_size = (224,224)
self.dynamic = 2
def __len__(self):
return 5000
def aug_data(self, image):
transform = A.Compose([
A.RandomBrightnessContrast(p=0.5),
])
transformed = transform(image=image.astype(np.uint8))
transformed_image = transformed["image"]
return transformed_image
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 and w > W:
pass_flag = False
elif mode == 'min':
if h < H and w < W:
pass_flag = False
return pass_flag
def get_sample(self, idx):
cloth_dir = os.path.join(self.image_root, self.data[idx])
ref_image_path = os.path.join(cloth_dir, 'target.jpg')
ref_image = cv2.imread(ref_image_path)
ref_image = cv2.cvtColor(ref_image.copy(), cv2.COLOR_BGR2RGB)
ref_mask_path = os.path.join(cloth_dir,'mask.jpg')
ref_mask = cv2.imread(ref_mask_path)[:,:,0] > 128
target_dirs = [i for i in os.listdir(cloth_dir ) if '.jpg' not in i]
target_dir_name = np.random.choice(target_dirs)
target_image_path = os.path.join(cloth_dir, target_dir_name + '.jpg')
target_image= cv2.imread(target_image_path)
tar_image = cv2.cvtColor(target_image.copy(), cv2.COLOR_BGR2RGB)
target_mask_path = os.path.join(cloth_dir, target_dir_name, 'segment.png')
tar_mask= cv2.imread(target_mask_path)[:,:,0]
target_mask = tar_mask == 7
kernel = np.ones((3, 3), dtype=np.uint8)
tar_mask = cv2.erode(target_mask.astype(np.uint8), kernel, iterations=3)
item_with_collage = self.process_pairs(ref_image, ref_mask, tar_image, tar_mask, max_ratio = 1.0)
sampled_time_steps = self.sample_timestep()
item_with_collage['time_steps'] = sampled_time_steps
return item_with_collage
|