scales annotations

This commit is contained in:
franksim
2025-11-09 13:01:46 +01:00
parent 85fd327052
commit 93d752209f
52 changed files with 84 additions and 2 deletions

View File

@@ -29,7 +29,7 @@ class MMP_Dataset(torch.utils.data.Dataset):
@param is_test: Whether this is the test set (True) or the validation/training set (False)
"""
self.image_size = image_size
self.images = []
self.images: Sequence[Tuple[str, Sequence[AnnotationRect]]] = []
self.anchor_grid = anchor_grid
self.min_iou = min_iou
self.is_test = is_test
@@ -72,8 +72,14 @@ class MMP_Dataset(torch.utils.data.Dataset):
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=self.images[idx][1], min_iou=self.min_iou
anchor_grid=self.anchor_grid, gts=scaled_annotations, min_iou=self.min_iou
)
return (img_tensor, label_grid, int(img_id))