198 lines
6.0 KiB
Python
198 lines
6.0 KiB
Python
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):
|
|
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 = []
|
|
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
|
|
"""
|
|
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:
|
|
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.
|
|
"""
|
|
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()
|