Spaces:
Runtime error
Runtime error
File size: 6,501 Bytes
7ae68fe |
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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
# --------------------------------------------------------
# InstructDiffusion
# Based on instruct-pix2pix (https://github.com/timothybrooks/instruct-pix2pix)
# Modified by Binxin Yang (tennyson@mail.ustc.edu.cn)
# --------------------------------------------------------
from __future__ import annotations
import json
import math
from pathlib import Path
from typing import Any
import numpy as np
import torch
import torchvision
from einops import rearrange
from PIL import Image
from torch.utils.data import Dataset
import cv2
import os
import random
import copy
from glob import glob
class COCOStuffDataset(Dataset):
def __init__(
self,
path: str,
path_edit: str = "None",
split: str = "train",
splits: tuple[float, float, float] = (0.9, 0.05, 0.05),
crop_res: int = 256,
flip_prob: float = 0.0,
transparency: float = 0,
batch_size: int = 10,
empty_percentage: float = 0,
):
assert split in ("train2017", "val2017")
assert sum(splits) == 1
self.split = split
self.path = path
self.path_edit = path_edit
self.batch_size = batch_size
self.crop_res = crop_res
self.flip_prob = flip_prob
self.empty_percentage = empty_percentage
self.transparency = transparency
if self.split in ["train2017", "val2017"]:
file_list = sorted(glob(os.path.join(self.path, "images", self.split, "*.jpg")))
assert len(file_list) > 0, "{} has no image".format(
os.path.join(self.path, "images", self.split)
)
file_list = [f.split("/")[-1].replace(".jpg", "") for f in file_list]
self.files = file_list
else:
raise ValueError("Invalid split name: {}".format(self.split))
seg_diverse_prompt_path = 'dataset/prompt/prompt_seg.txt'
self.seg_diverse_prompt_list=[]
with open(seg_diverse_prompt_path) as f:
line=f.readline()
while line:
line=line.strip('\n')
self.seg_diverse_prompt_list.append(line)
line=f.readline()
color_list_file_path='dataset/prompt/color_list_train_small.txt'
self.color_list=[]
with open(color_list_file_path) as f:
line = f.readline()
while line:
line_split = line.strip('\n').split(" ")
if len(line_split)>1:
temp = []
for i in range(4):
temp.append(line_split[i])
self.color_list.append(temp)
line = f.readline()
coco_label_list_path = self.path + '/labels.txt'
self.label_dict={}
with open(coco_label_list_path) as f:
line = f.readline()
while line:
line_split = line.strip('\n').split(": ")
self.label_dict[int(line_split[0])]=line_split[1]
line = f.readline()
def __len__(self) -> int:
length=len(self.files)
return length
def _augmentation_new(self, image, label):
# Cropping
h, w = label.shape
if h > w:
start_h = random.randint(0, h - w)
end_h = start_h + w
image = image[start_h:end_h]
label = label[start_h:end_h]
elif h < w:
start_w = random.randint(0, w - h)
end_w = start_w + h
image = image[:, start_w:end_w]
label = label[:, start_w:end_w]
else:
pass
image = Image.fromarray(image).resize((self.crop_res, self.crop_res), resample=Image.Resampling.LANCZOS)
image = np.asarray(image, dtype=np.uint8)
label = Image.fromarray(label).resize((self.crop_res, self.crop_res), resample=Image.Resampling.NEAREST)
label = np.asarray(label, dtype=np.int64)
return image, label
def __getitem__(self, i):
image_id = self.files[i]
img_path = os.path.join(self.path, "images", self.split, image_id + ".jpg")
mask_path = os.path.join(self.path, "annotations", self.split, image_id + ".png")
label = Image.open(mask_path).convert("L")
image = Image.open(img_path).convert("RGB")
label = np.asarray(label)
image = np.asarray(image)
image, label = self._augmentation_new(image,label)
label_list = np.unique(label)
label_list = list(label_list)
label_list_rest = [i for i in range(182)]
for item in label_list_rest:
if item in label_list:
label_list_rest.remove(item)
if 255 in label_list:
label_list.remove(255)
if len(label_list)!=0:
label_idx = random.choice(label_list)
if random.uniform(0, 1) < self.empty_percentage:
label_idx = random.choice(label_list_rest)
class_name = self.label_dict[label_idx+1]
prompt = random.choice(self.seg_diverse_prompt_list)
color = random.choice(self.color_list)
color_name = color[0]
prompt = prompt.format(color=color_name.lower(), object=class_name.lower())
R, G, B = color[3].split(",")
R = int(R)
G = int(G)
B = int(B)
else:
label_idx = 200
prompt = "leave the picture as it is."
mask = (label==label_idx)
image_0 = Image.fromarray(image)
image_1 = copy.deepcopy(image)
if len(label_list)!=0:
image_1[:,:,0][mask]=self.transparency*image_1[:,:,0][mask]+(1-self.transparency)*R
image_1[:,:,1][mask]=self.transparency*image_1[:,:,1][mask]+(1-self.transparency)*G
image_1[:,:,2][mask]=self.transparency*image_1[:,:,2][mask]+(1-self.transparency)*B
image_1 = Image.fromarray(image_1)
# return image_0, image_1, prompt
image_0 = rearrange(2 * torch.tensor(np.array(image_0)).float() / 255 - 1, "h w c -> c h w")
image_1 = rearrange(2 * torch.tensor(np.array(image_1)).float() / 255 - 1, "h w c -> c h w")
mask = torch.tensor(mask).float()
crop = torchvision.transforms.RandomCrop(self.crop_res)
flip = torchvision.transforms.RandomHorizontalFlip(float(self.flip_prob))
image_0, image_1 = flip(crop(torch.cat((image_0, image_1)))).chunk(2)
return dict(edited=image_1, edit=dict(c_concat=image_0, c_crossattn=prompt)) |