def __getitem__(self, item: int) -> Tuple[Tensor, Tensor, Tensor]:
"""
Gives the annotations of an image for all tasks. Also generates appropriate masks if the image
doesn't have annotations for a particular image.
:param item: the index of the image.
:return: tuple(x, t, m),
x is the transformed image data for different boxes, bim: [B x 3 x 224 x 224]
t is the target. t[i, tn] is 1 if the image i i first choice for task tn. Otherwise it is 0. [B x task_numbers]
m is the mask. m[tn] is 1 if this image is annotated for task tn. Otherwise it is 0. [task_numbers]
"""
# first choose MAX_GPU size number of annotations randomly.
the_image_id = self.all_image_ids[item]
valid_task_numbers = self.image_tasks[the_image_id]
some_coco = self.task_cocos[valid_task_numbers[0]]
the_image_dict = some_coco.loadImgs(the_image_id)[0]
I = load_image(get_image_file_name(the_image_dict))
all_image_anns = some_coco.loadAnns(
some_coco.getAnnIds(imgIds=the_image_id))
if len(all_image_anns) < MAX_GPU_SIZE:
selected_image_anns = all_image_anns
else:
selected_image_anns = np.random.choice(all_image_anns,
size=MAX_GPU_SIZE)
B = len(selected_image_anns)
x = [
image_transforms(
crop_img_to_bbox(I, target_transforms(a["bbox"], I.size)))
for a in selected_image_anns
]
t = np.zeros((B, len(TASK_NUMBERS)), dtype=np.float32)
m = np.zeros(len(TASK_NUMBERS), dtype=np.uint8)
for task_number in valid_task_numbers:
task_anns = self.task_cocos[task_number].loadAnns(
[a["id"] for a in selected_image_anns])
task_id = task_number_to_task_id(task_number)
for i, a in enumerate(task_anns):
if a["category_id"] == 1:
t[i, task_id] = 1
m[task_id] = 1
return default_collate(x), torch.tensor(t).float(), torch.tensor(
m).byte()
def __getitem__(self,
item) -> Tuple[Tensor, Tensor, Tensor, Tensor, List[dict]]:
"""
:param item:
:return: tuple (x, c, d, bbox, detections)
"""
the_image_id = self.list_of_valid_images[item]
the_image_dict = self.task_coco.loadImgs(the_image_id)[0]
related_annotations = self.task_coco.loadAnns(
self.task_coco.getAnnIds(imgIds=the_image_id))
detections = [{
"bbox": a["bbox"],
"score": 1.0,
"category_id": a["COCO_category_id"],
"image_id": the_image_id,
} for a in related_annotations]
if len(detections) > MAX_GPU_SIZE * 2:
print("WARNING, test image num detection larger than MAX_GPU_SIZE")
# TODO: fix this later
detections = detections[:MAX_GPU_SIZE * 2]
I = load_image(get_image_file_name(the_image_dict))
# fill in cropped images
x = [
image_transforms(
crop_img_to_bbox(I, target_transforms(det["bbox"], I.size)))
for det in detections
]
bbox = get_bbox_array_from_annotations(detections, I.size)
c = get_one_hot([(a["category_id"] - 1) for a in detections],
num_classes=90)
# fill in the detection scores
d = [float(det["score"]) for det in detections]
return (
default_collate(x),
torch.tensor(c).float(),
Tensor(d).unsqueeze(1),
torch.tensor(bbox).float(),
detections,
)
def __getitem__(self, item: int) -> Tuple[Tensor, Tensor, Tensor]:
"""
:return: (img_crop1, img_crop2, rank_posterior)
img_cropi: [3 x 224 x 224]
rank_posterior: []
"""
image_id, ann_id1, ann_id2, rankposterior = self.pairs[item]
the_image_dict = self.task_coco.loadImgs(image_id)[0]
the_ann1_dict = self.task_coco.loadAnns(ann_id1)[0]
the_ann2_dict = self.task_coco.loadAnns(ann_id2)[0]
# load image
I = load_image(get_image_file_name(the_image_dict))
# transform ann1 and ann2
img1 = image_transforms(
crop_img_to_bbox(I, target_transforms(the_ann1_dict["bbox"],
I.size)))
img2 = image_transforms(
crop_img_to_bbox(I, target_transforms(the_ann2_dict["bbox"],
I.size)))
return img1, img2, torch.tensor(rankposterior)
def __getitem__(self, item: int) -> Tuple[Tensor, dict]:
"""
:return: (img_crop, detection)
img_crop: [3 x 224 x 224]
detection: dict
"""
det = self.detections[item]
the_image_dict = self.task_coco.loadImgs(
self.task_coco.getImgIds(imgIds=det["image_id"]))[0]
img = load_image(get_image_file_name(the_image_dict))
img_crop = image_transforms(
crop_img_to_bbox(img, target_transforms(det["bbox"], img.size)))
detection = {
"bbox": det["bbox"],
"score": 1.0,
"category_id": det["category_id"],
"image_id": det["image_id"],
}
return img_crop, detection
def __getitem__(
self,
item: int) -> Tuple[Tensor, Tensor, Tensor, Tensor, List[dict]]:
"""
Give me an item. An item is basically the annotation of an image.
:param item: the index of the item to get
:return: tuple (x, c, d, t, detections),
x is the transformed image data for different boxes, dim: [B x 3 x 224x 224]
c is the one hot encoding of the classes, dim: [B x num_classes]
d is the detection scores for those boxes, dim: [B x 1]
t is the target which shows whether each of those boxes are first choice or not, dim: [B x 1]
detections is the json file containing the raw detections, type: List[dict]
"""
the_image_id = self.list_of_valid_images[item]
the_image_dict = self.task_coco.loadImgs(the_image_id)[0]
preferred_anns = self.task_coco.loadAnns(
self.task_coco.getAnnIds(imgIds=the_image_id, catIds=1))
non_preferred_anns = self.task_coco.loadAnns(
self.task_coco.getAnnIds(imgIds=the_image_id, catIds=0))
I = load_image(get_image_file_name(the_image_dict))
anns = []
t = []
# what should be the size?
size = int(round(np.random.poisson(self.len_lambda)))
size += 1 # ensure size is larger than 1
if size > MAX_GPU_SIZE: # ensure that size is smaller than MAX_GPU_SIZE
size = MAX_GPU_SIZE
# add from preferred_anns
number_of_preferred = len(preferred_anns)
if number_of_preferred > 0:
which_preferred_to_add = np.random.choice(preferred_anns,
size=number_of_preferred,
replace=False)
anns.extend([a for a in which_preferred_to_add])
t.extend([1] * number_of_preferred)
size -= number_of_preferred
if size > 0:
# add non preferred anns
number_of_non_preferred = len(non_preferred_anns)
if number_of_non_preferred > size:
number_of_non_preferred = size
# this is the maximum
number_of_non_preferred = min(number_of_non_preferred,
len(non_preferred_anns))
if number_of_non_preferred > 0:
which_non_preferred_to_add = np.random.choice(
non_preferred_anns,
size=number_of_non_preferred,
replace=False)
anns.extend([a for a in which_non_preferred_to_add])
t.extend([0] * number_of_non_preferred)
size -= number_of_non_preferred
# choose detections for annotation ids
detections = [{
"bbox": a["bbox"],
"score": 1.0,
"category_id": a["COCO_category_id"]
} for a in anns]
# fill in cropped images
x = [
image_transforms(
crop_img_to_bbox(I, target_transforms(det["bbox"], I.size)))
for det in detections
]
c = get_one_hot([(a["category_id"] - 1) for a in detections],
num_classes=90)
# fill in the detection scores
d = [float(det["score"]) for det in detections]
# unsqueeze to make [B] -> [B x 1]
return (
default_collate(x),
Tensor(c).float(),
Tensor(d).unsqueeze(1),
Tensor(t).unsqueeze(1),
detections,
)