small impovements

mmp/a4/dataset.py

@@ -5,7 +5,7 @@ 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
+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
@@ -33,6 +33,7 @@ class MMP_Dataset(torch.utils.data.Dataset):
         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))
 
@@ -126,15 +127,38 @@ def calculate_max_coverage(loader: DataLoader, min_iou: float) -> float:
     @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.
     """
-    raise NotImplementedError()
+    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 print_img_tensor_with_annotations(
-    img: torch.Tensor, annotations: Sequence["AnnotationRect"], output_file: str
+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).cpu().numpy()
-    img_np = img_np.astype(np.uint8)
+    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)
@@ -152,27 +176,30 @@ def print_img_tensor_with_annotations(
     plt.close(fig)
 
 
-def print_positive_boxes(
+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,
-    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(
+    draw_image_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",
+        output_file=f"mmp/a4/.output/{img_id}_transformed.png",
     )
 
 
@@ -195,12 +222,11 @@ def main():
     )
 
     for img, label, img_id in islice(dataloader, 12):
-        print_positive_boxes(
-            img_tensor=img[5],
+        draw_positive_boxes(
+            img_tensor=denormalize_image_tensor(img=img[5]),
             label_grid=label[5],
             img_id=img_id[5],
             anchor_grid=anchor_grid,
-            path_to_data=".data/mmp-public-3.2/train",
         )