| 11 | class CocoBase(Dataset): |
| 12 | """needed for (image, caption, segmentation) pairs""" |
| 13 | def __init__(self, size=None, dataroot="", datajson="", onehot_segmentation=False, use_stuffthing=False, |
| 14 | crop_size=None, force_no_crop=False, given_files=None, use_segmentation=True,crop_type=None): |
| 15 | self.split = self.get_split() |
| 16 | self.size = size |
| 17 | if crop_size is None: |
| 18 | self.crop_size = size |
| 19 | else: |
| 20 | self.crop_size = crop_size |
| 21 | |
| 22 | assert crop_type in [None, 'random', 'center'] |
| 23 | self.crop_type = crop_type |
| 24 | self.use_segmenation = use_segmentation |
| 25 | self.onehot = onehot_segmentation # return segmentation as rgb or one hot |
| 26 | self.stuffthing = use_stuffthing # include thing in segmentation |
| 27 | if self.onehot and not self.stuffthing: |
| 28 | raise NotImplemented("One hot mode is only supported for the " |
| 29 | "stuffthings version because labels are stored " |
| 30 | "a bit different.") |
| 31 | |
| 32 | data_json = datajson |
| 33 | with open(data_json) as json_file: |
| 34 | self.json_data = json.load(json_file) |
| 35 | self.img_id_to_captions = dict() |
| 36 | self.img_id_to_filepath = dict() |
| 37 | self.img_id_to_segmentation_filepath = dict() |
| 38 | |
| 39 | assert data_json.split("/")[-1] in [f"captions_train{self.year()}.json", |
| 40 | f"captions_val{self.year()}.json"] |
| 41 | # TODO currently hardcoded paths, would be better to follow logic in |
| 42 | # cocstuff pixelmaps |
| 43 | if self.use_segmenation: |
| 44 | if self.stuffthing: |
| 45 | self.segmentation_prefix = ( |
| 46 | f"data/cocostuffthings/val{self.year()}" if |
| 47 | data_json.endswith(f"captions_val{self.year()}.json") else |
| 48 | f"data/cocostuffthings/train{self.year()}") |
| 49 | else: |
| 50 | self.segmentation_prefix = ( |
| 51 | f"data/coco/annotations/stuff_val{self.year()}_pixelmaps" if |
| 52 | data_json.endswith(f"captions_val{self.year()}.json") else |
| 53 | f"data/coco/annotations/stuff_train{self.year()}_pixelmaps") |
| 54 | |
| 55 | imagedirs = self.json_data["images"] |
| 56 | self.labels = {"image_ids": list()} |
| 57 | for imgdir in tqdm(imagedirs, desc="ImgToPath"): |
| 58 | self.img_id_to_filepath[imgdir["id"]] = os.path.join(dataroot, imgdir["file_name"]) |
| 59 | self.img_id_to_captions[imgdir["id"]] = list() |
| 60 | pngfilename = imgdir["file_name"].replace("jpg", "png") |
| 61 | if self.use_segmenation: |
| 62 | self.img_id_to_segmentation_filepath[imgdir["id"]] = os.path.join( |
| 63 | self.segmentation_prefix, pngfilename) |
| 64 | if given_files is not None: |
| 65 | if pngfilename in given_files: |
| 66 | self.labels["image_ids"].append(imgdir["id"]) |
| 67 | else: |
| 68 | self.labels["image_ids"].append(imgdir["id"]) |
| 69 | |
| 70 | capdirs = self.json_data["annotations"] |