diff --git a/mmp/a4/2242167_original.png b/mmp/a4/2242167_original.png new file mode 100644 index 0000000..3b7270c Binary files /dev/null and b/mmp/a4/2242167_original.png differ diff --git a/mmp/a4/2242167_transformed.png b/mmp/a4/2242167_transformed.png new file mode 100644 index 0000000..f0ffc31 Binary files /dev/null and b/mmp/a4/2242167_transformed.png differ diff --git a/mmp/a4/2243763_original.png b/mmp/a4/2243763_original.png new file mode 100644 index 0000000..6cef7f2 Binary files /dev/null and b/mmp/a4/2243763_original.png differ diff --git a/mmp/a4/2243763_transformed.png b/mmp/a4/2243763_transformed.png new file mode 100644 index 0000000..06c44bb Binary files /dev/null and b/mmp/a4/2243763_transformed.png differ diff --git a/mmp/a4/2244338_original.png b/mmp/a4/2244338_original.png new file mode 100644 index 0000000..037e199 Binary files /dev/null and b/mmp/a4/2244338_original.png differ diff --git a/mmp/a4/2244338_transformed.png b/mmp/a4/2244338_transformed.png new file mode 100644 index 0000000..f129fdc Binary files /dev/null and b/mmp/a4/2244338_transformed.png differ diff --git a/mmp/a4/2245478_original.png b/mmp/a4/2245478_original.png new file mode 100644 index 0000000..7f8adcb Binary files /dev/null and b/mmp/a4/2245478_original.png differ diff --git a/mmp/a4/2245478_transformed.png b/mmp/a4/2245478_transformed.png new file mode 100644 index 0000000..e5516b1 Binary files /dev/null and b/mmp/a4/2245478_transformed.png differ diff --git a/mmp/a4/2247660_original.png b/mmp/a4/2247660_original.png new file mode 100644 index 0000000..3e153b1 Binary files /dev/null and b/mmp/a4/2247660_original.png differ diff --git a/mmp/a4/2247660_transformed.png b/mmp/a4/2247660_transformed.png new file mode 100644 index 0000000..dba52d0 Binary files /dev/null and b/mmp/a4/2247660_transformed.png differ diff --git a/mmp/a4/2249754_original.png b/mmp/a4/2249754_original.png new file mode 100644 index 0000000..4776e5e Binary files /dev/null and b/mmp/a4/2249754_original.png differ diff --git a/mmp/a4/2249754_transformed.png b/mmp/a4/2249754_transformed.png new file mode 100644 index 0000000..95e4be1 Binary files /dev/null and b/mmp/a4/2249754_transformed.png differ diff --git a/mmp/a4/2251077_original.png b/mmp/a4/2251077_original.png new file mode 100644 index 0000000..2becb15 Binary files /dev/null and b/mmp/a4/2251077_original.png differ diff --git a/mmp/a4/2251077_transformed.png b/mmp/a4/2251077_transformed.png new file mode 100644 index 0000000..9e460e3 Binary files /dev/null and b/mmp/a4/2251077_transformed.png differ diff --git a/mmp/a4/2253265_original.png b/mmp/a4/2253265_original.png new file mode 100644 index 0000000..1580633 Binary files /dev/null and b/mmp/a4/2253265_original.png differ diff --git a/mmp/a4/2253265_transformed.png b/mmp/a4/2253265_transformed.png new file mode 100644 index 0000000..c94fc86 Binary files /dev/null and b/mmp/a4/2253265_transformed.png differ diff --git a/mmp/a4/2256489_original.png b/mmp/a4/2256489_original.png new file mode 100644 index 0000000..ac4a202 Binary files /dev/null and b/mmp/a4/2256489_original.png differ diff --git a/mmp/a4/2256489_transformed.png b/mmp/a4/2256489_transformed.png new file mode 100644 index 0000000..ccf641b Binary files /dev/null and b/mmp/a4/2256489_transformed.png differ diff --git a/mmp/a4/2259929_original.png b/mmp/a4/2259929_original.png new file mode 100644 index 0000000..2bf598e Binary files /dev/null and b/mmp/a4/2259929_original.png differ diff --git a/mmp/a4/2259929_transformed.png b/mmp/a4/2259929_transformed.png new file mode 100644 index 0000000..1bdb869 Binary files /dev/null and b/mmp/a4/2259929_transformed.png differ diff --git a/mmp/a4/2263954_original.png b/mmp/a4/2263954_original.png new file mode 100644 index 0000000..09469e0 Binary files /dev/null and b/mmp/a4/2263954_original.png differ diff --git a/mmp/a4/2263954_transformed.png b/mmp/a4/2263954_transformed.png new file mode 100644 index 0000000..1af4fec Binary files /dev/null and b/mmp/a4/2263954_transformed.png differ diff --git a/mmp/a4/2266873_original.png b/mmp/a4/2266873_original.png new file mode 100644 index 0000000..8e68e87 Binary files /dev/null and b/mmp/a4/2266873_original.png differ diff --git a/mmp/a4/2266873_transformed.png b/mmp/a4/2266873_transformed.png new file mode 100644 index 0000000..f3f4775 Binary files /dev/null and b/mmp/a4/2266873_transformed.png differ diff --git a/mmp/a4/dataset.py b/mmp/a4/dataset.py index 14dce49..d642866 100644 --- a/mmp/a4/dataset.py +++ b/mmp/a4/dataset.py @@ -1,7 +1,17 @@ -from typing import Tuple +import os +import re +from typing import Sequence, Tuple import numpy as np import torch from torch.utils.data import DataLoader +from ..a3.annotation import read_groundtruth_file, AnnotationRect +from .label_grid import get_label_grid, draw_annotation_rects +import matplotlib.pyplot as plt +import matplotlib.patches as patches +from .anchor_grid import get_anchor_grid +from PIL import Image +from torchvision.transforms import transforms +from itertools import islice class MMP_Dataset(torch.utils.data.Dataset): @@ -18,16 +28,59 @@ class MMP_Dataset(torch.utils.data.Dataset): @param min_iou: The minimum IoU that is required for an overlap for the label grid. @param is_test: Whether this is the test set (True) or the validation/training set (False) """ - raise NotImplementedError() + self.image_size = image_size + self.images = [] + self.anchor_grid = anchor_grid + self.min_iou = min_iou + self.is_test = is_test + img_pattern = re.compile(r"^(\d+)\.jpg$") + files = set(os.listdir(path_to_data)) + + for fname in files: + match = img_pattern.match(fname) + if match: + img_file = os.path.join(path_to_data, fname) + annotations = read_groundtruth_file( + os.path.join(path_to_data, f"{match.group(1)}.gt_data.txt") + ) + self.images.append((img_file, annotations)) + + self.images.sort( + key=lambda x: int(re.match(r"(.*/)(\d+)(\.jpg)", x[0]).group(2)) + ) def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, int]: """ @return: 3-tuple of image tensor, label grid, and image (file-)number """ - raise NotImplementedError() + img = Image.open(self.images[idx][0]).convert("RGB") + padding = self.__padding__(img) + transform = transforms.Compose( + [ + transforms.Pad(padding, 0), + transforms.Resize((self.image_size, self.image_size)), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ), + ] + ) + img_tensor = transform(img) + label_grid = get_label_grid( + anchor_grid=self.anchor_grid, gts=self.images[idx][1], min_iou=self.min_iou + ) + img_id = re.match(r".*(\/)([0-9]+)(\.[^\/]*$)", self.images[idx][0]).group(2) + return (img_tensor, label_grid, int(img_id)) def __len__(self) -> int: - raise NotImplementedError() + return len(self.images) + + def __padding__(self, img) -> Tuple[int, int, int, int]: + w, h = img.size + size = max(w, h) + right_pad = size - w + bottom_pad = size - h + return (0, 0, right_pad, bottom_pad) def get_dataloader( @@ -37,8 +90,23 @@ def get_dataloader( num_workers: int, anchor_grid: np.ndarray, is_test: bool, + is_train: bool, ) -> DataLoader: - raise NotImplementedError() + dataset = MMP_Dataset( + path_to_data=path_to_data, + image_size=image_size, + is_test=is_test, + anchor_grid=anchor_grid, + min_iou=0.7, + ) + dataloader = DataLoader( + dataset, + batch_size=batch_size, + shuffle=is_train, + num_workers=num_workers, + pin_memory=True, + ) + return dataloader def calculate_max_coverage(loader: DataLoader, min_iou: float) -> float: @@ -48,3 +116,82 @@ def calculate_max_coverage(loader: DataLoader, min_iou: float) -> float: @return: Ratio of how mamy ground truth boxes are covered by a label grid box. Must be a value between 0 and 1. """ raise NotImplementedError() + + +def print_img_tensor_with_annotations( + img: torch.Tensor, annotations: Sequence["AnnotationRect"], output_file: str +): + # Convert tensor to numpy, permute dimensions + img_np = img.permute(1, 2, 0).cpu().numpy() + img_np = img_np.astype(np.uint8) + + fig, ax = plt.subplots(1) + ax.imshow(img_np) + for rect in annotations: + x1, y1, x2, y2 = rect.x1, rect.y1, rect.x2, rect.y2 + w = x2 - x1 + h = y2 - y1 + patch = patches.Rectangle( + (x1, y1), w, h, linewidth=2, edgecolor="red", facecolor="none" + ) + ax.add_patch(patch) + plt.axis("off") + plt.tight_layout(pad=0) + plt.savefig(output_file, bbox_inches="tight", pad_inches=0) + plt.close(fig) + + +def print_positive_boxes( + img_tensor: torch.Tensor, + label_grid: np.ndarray, + img_id: torch.Tensor, + anchor_grid: np.ndarray, + path_to_data: str, +): + annotations = [ + AnnotationRect.fromarray(anchor_grid[idx]) + for idx in np.ndindex(anchor_grid.shape[:-1]) + if label_grid[idx] + ] + print_img_tensor_with_annotations( + img_tensor, + annotations=annotations, + output_file=f"mmp/a4/{img_id}_transformed.png", + ) + draw_annotation_rects( + annotations=annotations, + image=f"{os.path.join(path_to_data, f'{str(img_id.item()).zfill(8)}.jpg')}", + output_path=f"mmp/a4/{img_id}_original.png", + ) + + +def main(): + anchor_grid = get_anchor_grid( + anchor_widths=[16, 32, 64, 96, 128, 144, 150, 160, 192, 224, 256], + aspect_ratios=[1 / 3, 1 / 2, 3 / 5, 2 / 3, 3 / 4, 1, 4 / 3, 5 / 3, 2, 2.5, 3], + num_rows=32, + num_cols=32, + scale_factor=20, + ) + dataloader = get_dataloader( + num_workers=6, + is_train=True, + is_test=False, + batch_size=8, + image_size=224, + path_to_data=".data/mmp-public-3.2/train", + anchor_grid=anchor_grid, + ) + + for img, label, img_id in islice(dataloader, 12): + print_positive_boxes( + img_tensor=img[5], + label_grid=label[5], + img_id=img_id[5], + anchor_grid=anchor_grid, + path_to_data=".data/mmp-public-3.2/train", + ) + + +if __name__ == "__main__": + main()