| | 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) |
| |
|
