diff --git a/pixano_inference/pytorch/deeplabv3.py b/pixano_inference/pytorch/deeplabv3.py index 050e126..b0bdef1 100644 --- a/pixano_inference/pytorch/deeplabv3.py +++ b/pixano_inference/pytorch/deeplabv3.py @@ -139,8 +139,7 @@ def preannotate( "mask": CompressedRLE.from_mask( unmold_mask(mask) ).to_dict(), - "category_id": int(label), - "category_name": voc_names(label), + "category": voc_names(label), } for label, mask in zip(labels, masks) ] diff --git a/pixano_inference/pytorch/maskrcnnv2.py b/pixano_inference/pytorch/maskrcnnv2.py index ed69b7e..810d0cc 100644 --- a/pixano_inference/pytorch/maskrcnnv2.py +++ b/pixano_inference/pytorch/maskrcnnv2.py @@ -140,8 +140,7 @@ def preannotate( "mask": CompressedRLE.from_mask( unmold_mask(output["masks"][i]) ).to_dict(), - "category_id": int(output["labels"][i]), - "category_name": coco_names_91(output["labels"][i]), + "category": coco_names_91(output["labels"][i]), } for i in range(len(output["scores"])) if output["scores"][i] > threshold diff --git a/pixano_inference/pytorch/yolov5.py b/pixano_inference/pytorch/yolov5.py index 8e6df2f..5347b6d 100644 --- a/pixano_inference/pytorch/yolov5.py +++ b/pixano_inference/pytorch/yolov5.py @@ -108,10 +108,7 @@ def preannotate( ) .normalize(h, w) .to_dict(), - "category_id": coco_ids_80to91(pred[5] + 1), - "category_name": coco_names_91( - coco_ids_80to91(pred[5] + 1) - ), + "category": coco_names_91(coco_ids_80to91(pred[5] + 1)), } for pred in img_output if pred[4] > threshold diff --git a/pixano_inference/tensorflow/efficientdet.py b/pixano_inference/tensorflow/efficientdet.py index 3cd8f24..a963245 100644 --- a/pixano_inference/tensorflow/efficientdet.py +++ b/pixano_inference/tensorflow/efficientdet.py @@ -104,8 +104,7 @@ def preannotate( ], confidence=float(output["detection_scores"][0][i]), ).to_dict(), - "category_id": int(output["detection_classes"][0][i]), - "category_name": coco_names_91( + "category": coco_names_91( output["detection_classes"][0][i] ), } diff --git a/pixano_inference/tensorflow/fasterrcnn.py b/pixano_inference/tensorflow/fasterrcnn.py index b9ab5e6..b892baf 100644 --- a/pixano_inference/tensorflow/fasterrcnn.py +++ b/pixano_inference/tensorflow/fasterrcnn.py @@ -106,8 +106,7 @@ def preannotate( ], confidence=float(output["detection_scores"][0][i]), ).to_dict(), - "category_id": int(output["detection_classes"][0][i]), - "category_name": coco_names_91( + "category": coco_names_91( output["detection_classes"][0][i] ), }