From aed21f1a1b62aeea1089694bee88c89368e26089 Mon Sep 17 00:00:00 2001 From: almaz Date: Fri, 30 Aug 2024 18:19:44 +0200 Subject: [PATCH] update geometry_to_task_type dict --- src/functions.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/functions.py b/src/functions.py index c72b6c0..e038260 100644 --- a/src/functions.py +++ b/src/functions.py @@ -1,11 +1,11 @@ import src.globals as g import supervisely as sly +from supervisely.nn import TaskType from supervisely.nn.inference import Session geometry_to_task_type = { - "object detection": [sly.Rectangle], - "instance segmentation": [sly.Bitmap, sly.Polygon, sly.AlphaMask], - # "semantic segmentation": [sly.Bitmap, sly.Polygon, sly.AlphaMask], + TaskType.OBJECT_DETECTION: [sly.Rectangle], + TaskType.INSTANCE_SEGMENTATION: [sly.Bitmap, sly.Polygon], } @@ -46,6 +46,7 @@ def get_classes(): return (matched_proj_cls, matched_model_cls), (not_matched_proj_cls, not_matched_model_cls) + def get_matched_class_names(): (cls_collection, _), _ = get_classes() - return [obj_cls.name for obj_cls in cls_collection] \ No newline at end of file + return [obj_cls.name for obj_cls in cls_collection]