implements dataset

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