| import os |
| import random |
|
|
| import torch |
| import torchvision.transforms as transforms |
| from PIL import Image |
|
|
|
|
| def recalculate_box_and_verify_if_valid(x, y, w, h, image_size, original_image_size, min_box_size): |
| scale = image_size / min(original_image_size) |
| crop_y = (original_image_size[1] * scale - image_size) // 2 |
| crop_x = (original_image_size[0] * scale - image_size) // 2 |
| x0 = max(x * scale - crop_x, 0) |
| y0 = max(y * scale - crop_y, 0) |
| x1 = min((x + w) * scale - crop_x, image_size) |
| y1 = min((y + h) * scale - crop_y, image_size) |
| if (x1 - x0) * (y1 - y0) / (image_size * image_size) < min_box_size: |
| return False, (None, None, None, None) |
| return True, (x0, y0, x1, y1) |
|
|
|
|
| class COCODataset(torch.utils.data.Dataset): |
| def __init__( |
| self, |
| data_path, |
| image_path, |
| image_size=512, |
| min_box_size=0.01, |
| max_boxes_per_data=8, |
| tokenizer=None, |
| ): |
| super().__init__() |
| self.min_box_size = min_box_size |
| self.max_boxes_per_data = max_boxes_per_data |
| self.image_size = image_size |
| self.image_path = image_path |
| self.tokenizer = tokenizer |
| self.transforms = transforms.Compose( |
| [ |
| transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BILINEAR), |
| transforms.CenterCrop(image_size), |
| transforms.ToTensor(), |
| transforms.Normalize([0.5], [0.5]), |
| ] |
| ) |
|
|
| self.data_list = torch.load(data_path, map_location="cpu") |
|
|
| def __getitem__(self, index): |
| if self.max_boxes_per_data > 99: |
| assert False, "Are you sure setting such large number of boxes per image?" |
|
|
| out = {} |
|
|
| data = self.data_list[index] |
| image = Image.open(os.path.join(self.image_path, data["file_path"])).convert("RGB") |
| original_image_size = image.size |
| out["pixel_values"] = self.transforms(image) |
|
|
| annos = data["annos"] |
|
|
| areas, valid_annos = [], [] |
| for anno in annos: |
| |
| x0, y0, x1, y1 = anno["bbox"] |
| x, y, w, h = x0, y0, x1 - x0, y1 - y0 |
| valid, (x0, y0, x1, y1) = recalculate_box_and_verify_if_valid( |
| x, y, w, h, self.image_size, original_image_size, self.min_box_size |
| ) |
| if valid: |
| anno["bbox"] = [x0, y0, x1, y1] |
| areas.append((x1 - x0) * (y1 - y0)) |
| valid_annos.append(anno) |
|
|
| |
| wanted_idxs = torch.tensor(areas).sort(descending=True)[1] |
| wanted_idxs = wanted_idxs[: self.max_boxes_per_data] |
| valid_annos = [valid_annos[i] for i in wanted_idxs] |
|
|
| out["boxes"] = torch.zeros(self.max_boxes_per_data, 4) |
| out["masks"] = torch.zeros(self.max_boxes_per_data) |
| out["text_embeddings_before_projection"] = torch.zeros(self.max_boxes_per_data, 768) |
|
|
| for i, anno in enumerate(valid_annos): |
| out["boxes"][i] = torch.tensor(anno["bbox"]) / self.image_size |
| out["masks"][i] = 1 |
| out["text_embeddings_before_projection"][i] = anno["text_embeddings_before_projection"] |
|
|
| prob_drop_boxes = 0.1 |
| if random.random() < prob_drop_boxes: |
| out["masks"][:] = 0 |
|
|
| caption = random.choice(data["captions"]) |
|
|
| prob_drop_captions = 0.5 |
| if random.random() < prob_drop_captions: |
| caption = "" |
| caption = self.tokenizer( |
| caption, |
| max_length=self.tokenizer.model_max_length, |
| padding="max_length", |
| truncation=True, |
| return_tensors="pt", |
| ) |
| out["caption"] = caption |
|
|
| return out |
|
|
| def __len__(self): |
| return len(self.data_list) |
|
|
