mmp_wise2526_franksim/mmp/a4/dataset.py

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, iou
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):
    def __init__(
        self,
        path_to_data: str,
        image_size: int,
        anchor_grid: np.ndarray,
        min_iou: float,
        is_test: bool,
    ):
        """
        @param anchor_grid: The anchor grid to be used for every image
        @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)
        """
        self.image_size = image_size
        self.images: Sequence[Tuple[str, Sequence[AnnotationRect]]] = []
        self.anchor_grid = anchor_grid
        self.min_iou = min_iou
        self.is_test = is_test
        self.path_to_data = path_to_data
        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)
                if is_test:
                    self.images.append((img_file, None))
                else:
                    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
        """
        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)
        img_id = re.match(r".*(\/)([0-9]+)(\.[^\/]*$)", self.images[idx][0]).group(2)
        if self.is_test:
            return (img_tensor, torch.Tensor(), int(img_id))

        scaled_annotations = []
        for annotation in self.images[idx][1]:
            annotation.scale(self.image_size / max(img.size[0], img.size[1]))
            scaled_annotations.append(annotation)

        label_grid = get_label_grid(
            anchor_grid=self.anchor_grid, gts=scaled_annotations, min_iou=self.min_iou
        )
        return (img_tensor, label_grid, int(img_id))

    def __len__(self) -> int:
        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(
    path_to_data: str,
    image_size: int,
    batch_size: int,
    num_workers: int,
    anchor_grid: np.ndarray,
    is_test: bool,
    is_train: bool = False,
) -> DataLoader:
    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:
    """
    @param loader: A DataLoader object, generated with the get_dataloader function.
    @param min_iou: Minimum IoU overlap that is required to count a ground truth box as covered.
    @return: Ratio of how mamy ground truth boxes are covered by a label grid box. Must be a value between 0 and 1.
    """
    total_boxes = 0
    covered_boxes = 0
    dataset: MMP_Dataset = loader.dataset
    anchor_grid = dataset.anchor_grid

    for img, _, img_id in islice(loader, 4):
        for batch_index in range(len(img)):
            gts_file = os.path.join(
                dataset.path_to_data,
                f"{str(img_id[batch_index].item()).zfill(8)}.gt_data.txt",
            )

            gts = read_groundtruth_file(gts_file)
            total_boxes += len(gts)
            for annotation in gts:
                for box_idx in np.ndindex(anchor_grid.shape[:-1]):
                    box_annotation = AnnotationRect.fromarray(anchor_grid[box_idx])
                    calculated_iou = iou(annotation, box_annotation)
                    if calculated_iou >= min_iou:
                        covered_boxes += 1
                        break
    return covered_boxes / total_boxes


def draw_image_tensor_with_annotations(
    img: torch.Tensor,
    annotations: Sequence["AnnotationRect"] | None,
    output_file: str,
):
    # Convert tensor to numpy, permute dimensions
    img_np = img.permute(1, 2, 0).numpy()
    img_np = np.clip(img_np, 0, 1)

    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 denormalize_image_tensor(
    img: torch.Tensor,
    mean=torch.tensor([0.485, 0.456, 0.406]).view(-1, 1, 1),
    std=torch.tensor([0.229, 0.224, 0.225]).view(-1, 1, 1),
) -> torch.Tensor:
    img_denormalized = img * std + mean
    return img_denormalized


def draw_positive_boxes(
    img_tensor: torch.Tensor,
    label_grid: np.ndarray,
    img_id: torch.Tensor,
    anchor_grid: np.ndarray,
):
    annotations = [
        AnnotationRect.fromarray(anchor_grid[idx])
        for idx in np.ndindex(anchor_grid.shape[:-1])
        if label_grid[idx]
    ]
    draw_image_tensor_with_annotations(
        img_tensor,
        annotations=annotations,
        output_file=f"mmp/a4/.output/{img_id}_transformed.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):
        draw_positive_boxes(
            img_tensor=denormalize_image_tensor(img=img[5]),
            label_grid=label[5],
            img_id=img_id[5],
            anchor_grid=anchor_grid,
        )


if __name__ == "__main__":
    main()