Spaces:
Runtime error
Runtime error
File size: 8,909 Bytes
0902a5f |
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 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 |
from typing import Dict
import numpy as np
from omegaconf import DictConfig, ListConfig
import torch
from torch.utils.data import Dataset
from pathlib import Path
import json
from PIL import Image
from torchvision import transforms
from einops import rearrange
from ldm.util import instantiate_from_config
# from datasets import load_dataset
import os
from collections import defaultdict
def make_multi_folder_data(paths, caption_files=None, **kwargs):
"""Make a concat dataset from multiple folders
Don't support captions yet
If paths is a list, that's ok, if it's a Dict interpret it as:
k=folder v=n_times to repeat that
"""
list_of_paths = []
if isinstance(paths, (Dict, DictConfig)):
assert caption_files is None, \
"Caption files not yet supported for repeats"
for folder_path, repeats in paths.items():
list_of_paths.extend([folder_path]*repeats)
paths = list_of_paths
if caption_files is not None:
datasets = [TextCapsDataset(p, caption_file=c, **kwargs) for (p, c) in zip(paths, caption_files)]
else:
datasets = [TextCapsDataset(p, **kwargs) for p in paths]
return torch.utils.data.ConcatDataset(datasets)
class TextCapsDataset(Dataset):
def __init__(self,
img_folder,
caption_file=None,
image_transforms=[],
first_stage_key = "jpg", cond_stage_key = "txt",
OneCapPerImage = False,
default_caption="",
ext="jpg",
postprocess=None,
return_paths=False,
filter_data=False,
filter_words=["sign", "poster"],
ocr_file=None,
) -> None:
"""Create a dataset from a folder of images.
If you pass in a root directory it will be searched for images
ending in ext (ext can be a list)
"""
self.root_dir = Path(img_folder)
self.first_stage_key = first_stage_key
self.cond_stage_key = cond_stage_key
# postprocess
if isinstance(postprocess, DictConfig):
postprocess = instantiate_from_config(postprocess)
self.postprocess = postprocess
# image transform
if isinstance(image_transforms, ListConfig):
image_transforms = [instantiate_from_config(tt) for tt in image_transforms]
image_transforms.extend([transforms.ToTensor(), # to be checked
transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
image_transforms = transforms.Compose(image_transforms)
self.tform = image_transforms
# caption
if caption_file is not None:
with open(caption_file, "rt") as f:
ext = Path(caption_file).suffix.lower()
if ext == ".json":
captions = json.load(f)
# elif ext == ".jsonl":
# lines = f.readlines()
# lines = [json.loads(x) for x in lines]
# captions = {x["file_name"]: x["text"].strip("\n") for x in lines}
else:
raise ValueError(f"Unrecognised format: {ext}")
self.captions = captions["data"]
if OneCapPerImage and ocr_file is None:
new_captions = []
taken_images = []
for caption_data in self.captions:
if caption_data["image_id"] in taken_images:
continue
else:
new_captions.append(caption_data)
taken_images.append(caption_data["image_id"])
self.captions = new_captions
else:
self.captions = None
if not isinstance(ext, (tuple, list, ListConfig)):
ext = [ext]
# Only used if there is no caption file
self.paths = []
for e in ext:
self.paths.extend(list(self.root_dir.rglob(f"*.{e}")))
self.default_caption = default_caption
self.return_paths = return_paths
self.filter_data = filter_data
self.filter_words = filter_words
self.ocr_file = ocr_file
self.ocr_data = []
if ocr_file is not None:
assert self.captions is not None
with open(ocr_file, "r") as f:
ocrs = json.loads(f.read())
ocr_data = ocrs['data']
self.ocr_data = ocr_data
def __len__(self):
if self.ocr_file is not None:
return len(self.ocr_data)
if self.captions is not None:
# return len(self.captions.keys())
return len(self.captions)
else:
return len(self.paths)
def __getitem__(self, index):
data = {}
if self.ocr_file is not None:
sample = self.ocr_data[index]
image_id = sample["image_id"]
ocr_tokens = sample["ocr_tokens"]
ocr_info = sample["ocr_info"]
chosen = image_id + ".jpg"
filename = self.root_dir/chosen
for d in self.captions:
if d["image_id"] == image_id:
image_captions = d["reference_strs"]
image_classes = d["image_classes"]
break
if not len(ocr_tokens) or not len(image_captions) or not len(image_classes):
return self.__getitem__(np.random.choice(self.__len__()))
tokens_state=defaultdict(list)
for token in ocr_tokens:
token_info = [
caption for caption in image_captions if (token.lower() in caption.rstrip(".").lower().split(" "))
]
tokens_state[len(token_info)].append(token.lower())
max_n = max(tokens_state.keys())
if max_n > 0:
valid_tokens = list(set(tokens_state[max_n]))
pos_info = dict()
for token in valid_tokens:
for item in ocr_info:
if item['word'].lower() == token:
token_box = item['bounding_box']
tx, ty = token_box['top_left_x'], token_box['top_left_y']
pos_info[token] = tx+ty
break
# arrange_tokens = list(dict(sorted(pos_info.items(), key=lambda x: x[1])).keys())
arrange_tokens = [item[0] for item in (sorted(pos_info.items(), key=lambda x: x[1]))]
valid_words = " ".join(arrange_tokens)
class_name = ""
for word in self.filter_words:
if word in " ".join(image_classes).lower():
class_name = word
break
if class_name == "":
return self.__getitem__(np.random.choice(self.__len__()))
else:
caption = "A {} that says '{}'.".format(
class_name, valid_words
)
else:
return self.__getitem__(np.random.choice(self.__len__()))
# if self.filter_data:
# if not len([word for word in self.filter_words if word in caption.rstrip(".").lower().split(" ")]):
# return self.__getitem__(np.random.choice(self.__len__()))
else:
if self.captions is not None:
# chosen = list(self.captions.keys())[index]
# caption = self.captions.get(chosen, None)
caption_data = self.captions[index]
chosen = os.path.basename(caption_data["image_path"])
caption = caption_data["caption_str"]
if caption is None:
caption = self.default_caption
filename = self.root_dir/chosen
# data[self.cond_stage_key] = caption
else:
filename = self.paths[index]
caption = self.default_caption
# data[self.cond_stage_key] = self.default_caption
if self.filter_data:
if not len([word for word in self.filter_words if word in caption.rstrip(".").lower().split(" ")]):
return self.__getitem__(np.random.choice(self.__len__()))
if self.return_paths:
data["path"] = str(filename)
im = Image.open(filename)
im = self.process_im(im)
data[self.first_stage_key] = im
data[self.cond_stage_key] = caption
# if self.captions is not None:
# data[self.cond_stage_key] = caption
# else:
# data[self.cond_stage_key] = self.default_caption
if self.postprocess is not None:
data = self.postprocess(data)
return data
def process_im(self, im):
im = im.convert("RGB")
return self.tform(im)
|