Spaces:
Sleeping
Sleeping
import os | |
import json | |
import albumentations | |
import numpy as np | |
from PIL import Image | |
from tqdm import tqdm | |
from torch.utils.data import Dataset | |
from abc import abstractmethod | |
class CocoBase(Dataset): | |
"""needed for (image, caption, segmentation) pairs""" | |
def __init__(self, size=None, dataroot="", datajson="", onehot_segmentation=False, use_stuffthing=False, | |
crop_size=None, force_no_crop=False, given_files=None, use_segmentation=True,crop_type=None): | |
self.split = self.get_split() | |
self.size = size | |
if crop_size is None: | |
self.crop_size = size | |
else: | |
self.crop_size = crop_size | |
assert crop_type in [None, 'random', 'center'] | |
self.crop_type = crop_type | |
self.use_segmenation = use_segmentation | |
self.onehot = onehot_segmentation # return segmentation as rgb or one hot | |
self.stuffthing = use_stuffthing # include thing in segmentation | |
if self.onehot and not self.stuffthing: | |
raise NotImplemented("One hot mode is only supported for the " | |
"stuffthings version because labels are stored " | |
"a bit different.") | |
data_json = datajson | |
with open(data_json) as json_file: | |
self.json_data = json.load(json_file) | |
self.img_id_to_captions = dict() | |
self.img_id_to_filepath = dict() | |
self.img_id_to_segmentation_filepath = dict() | |
assert data_json.split("/")[-1] in [f"captions_train{self.year()}.json", | |
f"captions_val{self.year()}.json"] | |
# TODO currently hardcoded paths, would be better to follow logic in | |
# cocstuff pixelmaps | |
if self.use_segmenation: | |
if self.stuffthing: | |
self.segmentation_prefix = ( | |
f"data/cocostuffthings/val{self.year()}" if | |
data_json.endswith(f"captions_val{self.year()}.json") else | |
f"data/cocostuffthings/train{self.year()}") | |
else: | |
self.segmentation_prefix = ( | |
f"data/coco/annotations/stuff_val{self.year()}_pixelmaps" if | |
data_json.endswith(f"captions_val{self.year()}.json") else | |
f"data/coco/annotations/stuff_train{self.year()}_pixelmaps") | |
imagedirs = self.json_data["images"] | |
self.labels = {"image_ids": list()} | |
for imgdir in tqdm(imagedirs, desc="ImgToPath"): | |
self.img_id_to_filepath[imgdir["id"]] = os.path.join(dataroot, imgdir["file_name"]) | |
self.img_id_to_captions[imgdir["id"]] = list() | |
pngfilename = imgdir["file_name"].replace("jpg", "png") | |
if self.use_segmenation: | |
self.img_id_to_segmentation_filepath[imgdir["id"]] = os.path.join( | |
self.segmentation_prefix, pngfilename) | |
if given_files is not None: | |
if pngfilename in given_files: | |
self.labels["image_ids"].append(imgdir["id"]) | |
else: | |
self.labels["image_ids"].append(imgdir["id"]) | |
capdirs = self.json_data["annotations"] | |
for capdir in tqdm(capdirs, desc="ImgToCaptions"): | |
# there are in average 5 captions per image | |
#self.img_id_to_captions[capdir["image_id"]].append(np.array([capdir["caption"]])) | |
self.img_id_to_captions[capdir["image_id"]].append(capdir["caption"]) | |
self.rescaler = albumentations.SmallestMaxSize(max_size=self.size) | |
if self.split=="validation": | |
self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size) | |
else: | |
# default option for train is random crop | |
if self.crop_type in [None, 'random']: | |
self.cropper = albumentations.RandomCrop(height=self.crop_size, width=self.crop_size) | |
else: | |
self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size) | |
self.preprocessor = albumentations.Compose( | |
[self.rescaler, self.cropper], | |
additional_targets={"segmentation": "image"}) | |
if force_no_crop: | |
self.rescaler = albumentations.Resize(height=self.size, width=self.size) | |
self.preprocessor = albumentations.Compose( | |
[self.rescaler], | |
additional_targets={"segmentation": "image"}) | |
def year(self): | |
raise NotImplementedError() | |
def __len__(self): | |
return len(self.labels["image_ids"]) | |
def preprocess_image(self, image_path, segmentation_path=None): | |
image = Image.open(image_path) | |
if not image.mode == "RGB": | |
image = image.convert("RGB") | |
image = np.array(image).astype(np.uint8) | |
if segmentation_path: | |
segmentation = Image.open(segmentation_path) | |
if not self.onehot and not segmentation.mode == "RGB": | |
segmentation = segmentation.convert("RGB") | |
segmentation = np.array(segmentation).astype(np.uint8) | |
if self.onehot: | |
assert self.stuffthing | |
# stored in caffe format: unlabeled==255. stuff and thing from | |
# 0-181. to be compatible with the labels in | |
# https://github.com/nightrome/cocostuff/blob/master/labels.txt | |
# we shift stuffthing one to the right and put unlabeled in zero | |
# as long as segmentation is uint8 shifting to right handles the | |
# latter too | |
assert segmentation.dtype == np.uint8 | |
segmentation = segmentation + 1 | |
processed = self.preprocessor(image=image, segmentation=segmentation) | |
image, segmentation = processed["image"], processed["segmentation"] | |
else: | |
image = self.preprocessor(image=image,)['image'] | |
image = (image / 127.5 - 1.0).astype(np.float32) | |
if segmentation_path: | |
if self.onehot: | |
assert segmentation.dtype == np.uint8 | |
# make it one hot | |
n_labels = 183 | |
flatseg = np.ravel(segmentation) | |
onehot = np.zeros((flatseg.size, n_labels), dtype=np.bool) | |
onehot[np.arange(flatseg.size), flatseg] = True | |
onehot = onehot.reshape(segmentation.shape + (n_labels,)).astype(int) | |
segmentation = onehot | |
else: | |
segmentation = (segmentation / 127.5 - 1.0).astype(np.float32) | |
return image, segmentation | |
else: | |
return image | |
def __getitem__(self, i): | |
img_path = self.img_id_to_filepath[self.labels["image_ids"][i]] | |
if self.use_segmenation: | |
seg_path = self.img_id_to_segmentation_filepath[self.labels["image_ids"][i]] | |
image, segmentation = self.preprocess_image(img_path, seg_path) | |
else: | |
image = self.preprocess_image(img_path) | |
captions = self.img_id_to_captions[self.labels["image_ids"][i]] | |
# randomly draw one of all available captions per image | |
caption = captions[np.random.randint(0, len(captions))] | |
example = {"image": image, | |
#"caption": [str(caption[0])], | |
"caption": caption, | |
"img_path": img_path, | |
"filename_": img_path.split(os.sep)[-1] | |
} | |
if self.use_segmenation: | |
example.update({"seg_path": seg_path, 'segmentation': segmentation}) | |
return example | |
class CocoImagesAndCaptionsTrain2017(CocoBase): | |
"""returns a pair of (image, caption)""" | |
def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,): | |
super().__init__(size=size, | |
dataroot="data/coco/train2017", | |
datajson="data/coco/annotations/captions_train2017.json", | |
onehot_segmentation=onehot_segmentation, | |
use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop) | |
def get_split(self): | |
return "train" | |
def year(self): | |
return '2017' | |
class CocoImagesAndCaptionsValidation2017(CocoBase): | |
"""returns a pair of (image, caption)""" | |
def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False, | |
given_files=None): | |
super().__init__(size=size, | |
dataroot="data/coco/val2017", | |
datajson="data/coco/annotations/captions_val2017.json", | |
onehot_segmentation=onehot_segmentation, | |
use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop, | |
given_files=given_files) | |
def get_split(self): | |
return "validation" | |
def year(self): | |
return '2017' | |
class CocoImagesAndCaptionsTrain2014(CocoBase): | |
"""returns a pair of (image, caption)""" | |
def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,crop_type='random'): | |
super().__init__(size=size, | |
dataroot="data/coco/train2014", | |
datajson="data/coco/annotations2014/annotations/captions_train2014.json", | |
onehot_segmentation=onehot_segmentation, | |
use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop, | |
use_segmentation=False, | |
crop_type=crop_type) | |
def get_split(self): | |
return "train" | |
def year(self): | |
return '2014' | |
class CocoImagesAndCaptionsValidation2014(CocoBase): | |
"""returns a pair of (image, caption)""" | |
def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False, | |
given_files=None,crop_type='center',**kwargs): | |
super().__init__(size=size, | |
dataroot="data/coco/val2014", | |
datajson="data/coco/annotations2014/annotations/captions_val2014.json", | |
onehot_segmentation=onehot_segmentation, | |
use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop, | |
given_files=given_files, | |
use_segmentation=False, | |
crop_type=crop_type) | |
def get_split(self): | |
return "validation" | |
def year(self): | |
return '2014' | |
if __name__ == '__main__': | |
with open("data/coco/annotations2014/annotations/captions_val2014.json", "r") as json_file: | |
json_data = json.load(json_file) | |
capdirs = json_data["annotations"] | |
import pudb; pudb.set_trace() | |
#d2 = CocoImagesAndCaptionsTrain2014(size=256) | |
d2 = CocoImagesAndCaptionsValidation2014(size=256) | |
print("constructed dataset.") | |
print(f"length of {d2.__class__.__name__}: {len(d2)}") | |
ex2 = d2[0] | |
# ex3 = d3[0] | |
# print(ex1["image"].shape) | |
print(ex2["image"].shape) | |
# print(ex3["image"].shape) | |
# print(ex1["segmentation"].shape) | |
print(ex2["caption"].__class__.__name__) | |