Spaces:
Build error
Build error
File size: 1,569 Bytes
169e11c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import numpy as np
from src.detection.vision.coco_utils import ConvertCocoPolysToMask, CocoDetection, _coco_remove_images_without_annotations
from src.detection.vision.transforms import Compose
class UpdateIsCrowd(object):
def __init__(self, min_size, min_vis=0.2):
self.min_size = min_size
self.min_vis = min_vis
def __call__(self, image, target):
for i, ann in enumerate(target['annotations']):
bbox = ann['bbox']
bbox_too_small = max(bbox[-1], bbox[-2]) < self.min_size
if 'vis' in ann:
vis = ann['vis']
elif 'keypoints' in ann:
vis = (np.array(ann['keypoints'])[2::3] == 2).mean().round(2)
else:
raise RuntimeError(
"The given annotations have no visibility measure. Are you sure you want to proceed?")
not_vis = vis < self.min_vis
target['annotations'][i]['iscrowd'] = max(
ann['iscrowd'], int(bbox_too_small), int(not_vis))
return image, target
def get_mot_dataset(img_folder, ann_file, transforms, min_size=25, min_vis=0.2):
t = [UpdateIsCrowd(min_size=min_size, min_vis=min_vis),
ConvertCocoPolysToMask()]
if transforms is not None:
t.append(transforms)
transforms = Compose(t)
dataset = CocoDetection(img_folder=img_folder,
ann_file=ann_file,
transforms=transforms)
dataset = _coco_remove_images_without_annotations(dataset)
return dataset
|