Formats image annotations for object detection, instance segmentation, and pose estimation tasks. The class standardizes the image and instance annotations to be used by the `collate_fn` in PyTorch DataLoader. Attributes: bbox_format (str): Format for bounding boxes. Default is
| 832 | |
| 833 | # TODO: technically this is not an augmentation, maybe we should put this to another files |
| 834 | class Format: |
| 835 | """ |
| 836 | Formats image annotations for object detection, instance segmentation, and pose estimation tasks. The class |
| 837 | standardizes the image and instance annotations to be used by the `collate_fn` in PyTorch DataLoader. |
| 838 | |
| 839 | Attributes: |
| 840 | bbox_format (str): Format for bounding boxes. Default is 'xywh'. |
| 841 | normalize (bool): Whether to normalize bounding boxes. Default is True. |
| 842 | return_mask (bool): Return instance masks for segmentation. Default is False. |
| 843 | return_keypoint (bool): Return keypoints for pose estimation. Default is False. |
| 844 | mask_ratio (int): Downsample ratio for masks. Default is 4. |
| 845 | mask_overlap (bool): Whether to overlap masks. Default is True. |
| 846 | batch_idx (bool): Keep batch indexes. Default is True. |
| 847 | """ |
| 848 | |
| 849 | def __init__(self, |
| 850 | bbox_format='xywh', |
| 851 | normalize=True, |
| 852 | return_mask=False, |
| 853 | return_keypoint=False, |
| 854 | mask_ratio=4, |
| 855 | mask_overlap=True, |
| 856 | batch_idx=True): |
| 857 | """Initializes the Format class with given parameters.""" |
| 858 | self.bbox_format = bbox_format |
| 859 | self.normalize = normalize |
| 860 | self.return_mask = return_mask # set False when training detection only |
| 861 | self.return_keypoint = return_keypoint |
| 862 | self.mask_ratio = mask_ratio |
| 863 | self.mask_overlap = mask_overlap |
| 864 | self.batch_idx = batch_idx # keep the batch indexes |
| 865 | |
| 866 | def __call__(self, labels): |
| 867 | """Return formatted image, classes, bounding boxes & keypoints to be used by 'collate_fn'.""" |
| 868 | img = labels.pop('img') |
| 869 | h, w = img.shape[:2] |
| 870 | cls = labels.pop('cls') |
| 871 | instances = labels.pop('instances') |
| 872 | instances.convert_bbox(format=self.bbox_format) |
| 873 | instances.denormalize(w, h) |
| 874 | nl = len(instances) |
| 875 | |
| 876 | if self.return_mask: |
| 877 | if nl: |
| 878 | masks, instances, cls = self._format_segments(instances, cls, w, h) |
| 879 | masks = torch.from_numpy(masks) |
| 880 | else: |
| 881 | masks = torch.zeros(1 if self.mask_overlap else nl, img.shape[0] // self.mask_ratio, |
| 882 | img.shape[1] // self.mask_ratio) |
| 883 | labels['masks'] = masks |
| 884 | if self.normalize: |
| 885 | instances.normalize(w, h) |
| 886 | labels['img'] = self._format_img(img) |
| 887 | labels['cls'] = torch.from_numpy(cls) if nl else torch.zeros(nl) |
| 888 | labels['bboxes'] = torch.from_numpy(instances.bboxes) if nl else torch.zeros((nl, 4)) |
| 889 | if self.return_keypoint: |
| 890 | labels['keypoints'] = torch.from_numpy(instances.keypoints) |
| 891 | # Then we can use collate_fn |
no outgoing calls
no test coverage detected