Spaces:
Runtime error
Runtime error
File size: 2,834 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 |
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
class DreamBoothDataset(BaseDataset):
def __init__(self, fg_dir, bg_dir):
self.bg_dir = bg_dir
bg_data = os.listdir(self.bg_dir)
self.bg_data = [i for i in bg_data if 'mask' in i]
self.image_dir = fg_dir
self.data = os.listdir(self.image_dir)
self.size = (512,512)
self.clip_size = (224,224)
'''
Dynamic:
0: Static View, High Quality
1: Multi-view, Low Quality
2: Multi-view, High Quality
'''
self.dynamic = 1
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
idx = np.random.randint(0, len(self.data)-1)
item = self.get_sample(idx)
return item
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_alpha_mask(self, mask_path):
image = cv2.imread( mask_path, cv2.IMREAD_UNCHANGED)
mask = (image[:,:,-1] > 128).astype(np.uint8)
return mask
def get_sample(self, idx):
dir_name = self.data[idx]
dir_path = os.path.join(self.image_dir, dir_name)
images = os.listdir(dir_path)
image_name = [i for i in images if '.png' in i][0]
image_path = os.path.join(dir_path, image_name)
image = cv2.imread( image_path, cv2.IMREAD_UNCHANGED)
mask = (image[:,:,-1] > 128).astype(np.uint8)
image = image[:,:,:-1]
image = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2RGB)
ref_image = image
ref_mask = mask
ref_image, ref_mask = expand_image_mask(image, mask, ratio=1.4)
bg_idx = np.random.randint(0, len(self.bg_data)-1)
tar_mask_name = self.bg_data[bg_idx]
tar_mask_path = os.path.join(self.bg_dir, tar_mask_name)
tar_image_path = tar_mask_path.replace('_mask','_GT')
tar_image = cv2.imread(tar_image_path).astype(np.uint8)
tar_image = cv2.cvtColor(tar_image, cv2.COLOR_BGR2RGB)
tar_mask = (cv2.imread(tar_mask_path) > 128).astype(np.uint8)[:,:,0]
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
|