xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor class_id assignment to Detections in VOC

This commit is contained in:
kirilllzaitsev 2023-07-29 00:39:11 +02:00
parent 446ae23d99
commit 9a9d0c5259
1 changed files with 6 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
supervision/dataset/formats/pascal_voc.py
View File

@ -195,24 +195,18 @@ def load_pascal_voc_annotations(
masks.append(mask_from_polygon)
xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
for k in set(class_names):
if k not in classes:
classes.append(k)
class_id = np.array([classes.index(class_name) for class_name in class_names])
if with_masks:
masks = np.array(masks)
annotation = Detections(
xyxy=xyxy, mask=masks, class_id=np.array(class_names)
)
annotation = Detections(xyxy=xyxy, mask=np.array(masks), class_id=class_id)
else:
annotation = Detections(xyxy=xyxy, class_id=np.array(class_names))
annotation = Detections(xyxy=xyxy, class_id=class_id)
images[image_path.name] = image
annotations[image_path.name] = annotation
classes += class_names
classes = list(set(classes))
for annotation in annotations.values():
class_id = [classes.index(class_name) for class_name in annotation.class_id]
annotation.class_id = np.array(class_id)
return classes, images, annotations