Spaces:
Moustached
/
Runtime error

File size: 7,917 Bytes
f9a674e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
import json, os, random, math
from collections import defaultdict
from copy import deepcopy

import torch
from torch.utils.data import Dataset
import torchvision.transforms as transforms

import numpy as np
from PIL import Image
from .base_dataset import BaseDataset, check_filenames_in_zipdata, recalculate_box_and_verify_if_valid 
from io import BytesIO



def not_in_at_all(list1, list2):
	for a in list1:
		if a in list2:
			return False
	return True


def clean_annotations(annotations):
	for anno in annotations:
		anno.pop("segmentation", None)
		anno.pop("area", None)
		anno.pop("iscrowd", None)
		# anno.pop("id", None)


def make_a_sentence(obj_names, clean=False):

	if clean:
		obj_names = [ name[:-6] if ("-other" in name) else name for name in obj_names]

	caption = ""
	tokens_positive = []
	for obj_name in obj_names:
		start_len = len(caption)
		caption += obj_name
		end_len = len(caption)
		caption += ", "
		tokens_positive.append(
			[[start_len, end_len]] # in real caption, positive tokens can be disjoint, thus using list of list
		)
	caption = caption[:-2] # remove last ", "

	return caption #, tokens_positive


def check_all_have_same_images(instances_data, stuff_data, caption_data):
	if stuff_data is not None:
		assert instances_data["images"] == stuff_data["images"]
	if caption_data is not None:
		assert instances_data["images"] == caption_data["images"]


class CDDataset(BaseDataset):
	"CD: Caption Detection"
	def __init__(self, 
                image_root,
				category_embedding_path,
				instances_json_path = None,
				stuff_json_path = None,
				caption_json_path = None,
				prob_real_caption = 0,
				fake_caption_type = 'empty',
                image_size=256, 
                max_images=None,
                min_box_size=0.01,
                max_boxes_per_image=8,
                include_other=False, 
				random_crop = False,
				random_flip = True,
                ):
		super().__init__(random_crop, random_flip, image_size)

		self.image_root = image_root
		self.category_embedding_path = category_embedding_path
		self.instances_json_path = instances_json_path
		self.stuff_json_path = stuff_json_path
		self.caption_json_path = caption_json_path
		self.prob_real_caption = prob_real_caption
		self.fake_caption_type = fake_caption_type
		self.max_images = max_images
		self.min_box_size = min_box_size
		self.max_boxes_per_image = max_boxes_per_image
		self.include_other = include_other


		assert fake_caption_type in ["empty", "made"]
		if prob_real_caption > 0:
			assert caption_json_path is not None, "caption json must be given"
	

		# Load all jsons 
		with open(instances_json_path, 'r') as f:
			instances_data = json.load(f) # keys: 'info', 'images', 'licenses', 'categories', 'annotations'
		clean_annotations(instances_data["annotations"])
		self.instances_data = instances_data

		self.stuff_data = None
		if stuff_json_path is not None:
			with open(stuff_json_path, 'r') as f:
				stuff_data = json.load(f) # keys: 'info', 'images', 'licenses', 'categories', 'annotations'
			clean_annotations(stuff_data["annotations"])
			self.stuff_data = stuff_data

		self.captions_data = None
		if caption_json_path is not None:
			with open(caption_json_path, 'r') as f:
				captions_data = json.load(f) # keys: 'info', 'images', 'licenses', 'categories', 'annotations'
			clean_annotations(captions_data["annotations"])
			self.captions_data = captions_data


		# Load preprocessed name embedding 
		self.category_embeddings = torch.load(category_embedding_path)
		self.embedding_len = list( self.category_embeddings.values() )[0].shape[0]
		

		# Misc  
		self.image_ids = [] # main list for selecting images
		self.image_id_to_filename = {} # file names used to read image
		check_all_have_same_images(self.instances_data, self.stuff_data, self.captions_data)
		for image_data in self.instances_data['images']:
			image_id = image_data['id']
			filename = image_data['file_name']
			self.image_ids.append(image_id)
			self.image_id_to_filename[image_id] = filename

		
		# All category names (including things and stuff)
		self.object_idx_to_name = {} 
		for category_data in self.instances_data['categories']:
			self.object_idx_to_name[category_data['id']] = category_data['name']
		if self.stuff_data is not None:
			for category_data in self.stuff_data['categories']:
				self.object_idx_to_name[category_data['id']] = category_data['name']


		# Add object data from instances and stuff 
		self.image_id_to_objects = defaultdict(list)
		self.select_objects( self.instances_data['annotations'] )
		if self.stuff_data is not None:
			self.select_objects( self.stuff_data['annotations'] )

		# Add caption data 
		if self.captions_data is not None:
			self.image_id_to_captions = defaultdict(list)
			self.select_captions( self.captions_data['annotations'] )

		# Check if all filenames can be found in the zip file
		# all_filenames = [self.image_id_to_filename[idx] for idx in self.image_ids]
		# check_filenames_in_zipdata(all_filenames, image_root)
		

	def select_objects(self, annotations):
		for object_anno in annotations:
			image_id = object_anno['image_id']
			object_name = self.object_idx_to_name[object_anno['category_id']]
			other_ok = object_name != 'other' or self.include_other
			if other_ok:
				self.image_id_to_objects[image_id].append(object_anno)


	def select_captions(self, annotations):
		for caption_data in annotations:
			image_id = caption_data['image_id']
			self.image_id_to_captions[image_id].append(caption_data)


	def total_images(self):
		return len(self)


	def __getitem__(self, index):
		if self.max_boxes_per_image > 99:
			assert False, "Are you sure setting such large number of boxes?"

		out = {}

		image_id = self.image_ids[index]
		out['id'] = image_id
		
		# Image 
		filename = self.image_id_to_filename[image_id]
		image = self.fetch_image(filename)
		#WW, HH = image.size
		image_tensor, trans_info = self.transform_image(image)
		out["image"] = image_tensor
	

		# Select valid boxes after cropping (center or random)
		this_image_obj_annos = deepcopy(self.image_id_to_objects[image_id])
		areas = []
		all_obj_names = []
		all_boxes = []
		all_masks = []
		all_positive_embeddings = []
		for object_anno in this_image_obj_annos:

			x, y, w, h = object_anno['bbox']
			valid, (x0, y0, x1, y1) = recalculate_box_and_verify_if_valid(x, y, w, h, trans_info, self.image_size, self.min_box_size)

			if valid:
				areas.append(  (x1-x0)*(y1-y0) )
				obj_name = self.object_idx_to_name[ object_anno['category_id']  ]
				all_obj_names.append(obj_name)
				all_boxes.append( torch.tensor([x0,y0,x1,y1]) / self.image_size ) # scale to 0-1
				all_masks.append(1)
				all_positive_embeddings.append( self.category_embeddings[obj_name]  )

		wanted_idxs = torch.tensor(areas).sort(descending=True)[1]
		wanted_idxs = wanted_idxs[0:self.max_boxes_per_image]
		obj_names = [] # used for making a sentence
		boxes = torch.zeros(self.max_boxes_per_image, 4)
		masks = torch.zeros(self.max_boxes_per_image)
		positive_embeddings = torch.zeros(self.max_boxes_per_image, self.embedding_len)
		for i, idx in enumerate(wanted_idxs):
			obj_names.append(  all_obj_names[idx]   )
			boxes[i] = all_boxes[idx]
			masks[i] = all_masks[idx]
			positive_embeddings[i] = all_positive_embeddings[idx]

		# Caption
		if random.uniform(0, 1) < self.prob_real_caption:
			caption_data = self.image_id_to_captions[image_id]
			idx = random.randint(0,  len(caption_data)-1 )
			caption = caption_data[idx]["caption"]
		else:
			if self.fake_caption_type == "empty":
				caption = ""
			else:
				caption = make_a_sentence(obj_names, clean=True)
		
		
		out["caption"] = caption
		out["boxes"] = boxes
		out["masks"] = masks
		out["positive_embeddings"] = positive_embeddings

		return out 


	def __len__(self):
		if self.max_images is None:
			return len(self.image_ids)
		return min(len(self.image_ids), self.max_images)