scales annotations
This commit is contained in:
@@ -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))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user