fix is_coco_dataset judgment condition

mindspore-lab · Dec 20, 2024 · 1590d55 · 1590d55
1 parent 4c563a1
commit 1590d55
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 5 deletions.
diff --git a/deploy/mslite_predict.py b/deploy/mslite_predict.py
@@ -160,7 +160,12 @@ def infer(args):
         img = cv2.imread(args.image_path)
     else:
         raise ValueError("Detect: input image file not available.")
-    is_coco_dataset = "coco" in args.data.dataset_name
+    # referred from https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/detect/val.py#L74
+    is_coco_dataset = (
+            isinstance(args.data.val_set, str)
+            and "coco" in args.data.val_set
+            and (args.data.val_set.endswith(f"{os.sep}val2017.txt") or args.data.val_set.endswith(f"{os.sep}test-dev2017.txt"))
+    )  # is COCO
     # Detect
     result_dict = detect(
         mindir_path=args.mindir_path,

diff --git a/examples/finetune_carparts_seg/predict.py b/examples/finetune_carparts_seg/predict.py
@@ -297,7 +297,12 @@ def infer(args):
         raise ValueError("Detect: input image file not available.")
 
     # Detect
-    is_coco_dataset = "coco" in args.data.dataset_name
+    # referred from https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/detect/val.py#L74
+    is_coco_dataset = (
+        isinstance(args.data.val_set, str)
+        and "coco" in args.data.val_set
+        and (args.data.val_set.endswith(f"{os.sep}val2017.txt") or args.data.val_set.endswith(f"{os.sep}test-dev2017.txt"))
+    )  # is COCO
     if args.task == "detect":
         result_dict = detect(
             network=network,

diff --git a/examples/finetune_visdrone/predict.py b/examples/finetune_visdrone/predict.py
@@ -297,7 +297,12 @@ def infer(args):
         raise ValueError("Detect: input image file not available.")
 
     # Detect
-    is_coco_dataset = "coco" in args.data.dataset_name
+    # referred from https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/detect/val.py#L74
+    is_coco_dataset = (
+            isinstance(args.data.val_set, str)
+            and "coco" in args.data.val_set
+            and (args.data.val_set.endswith(f"{os.sep}val2017.txt") or args.data.val_set.endswith(f"{os.sep}test-dev2017.txt"))
+    )  # is COCO
     if args.task == "detect":
         result_dict = detect(
             network=network,

diff --git a/test.py b/test.py
@@ -466,7 +466,12 @@ def main(args):
 
     # Create Dataloader
     dataset_path = args.data.val_set
-    is_coco_dataset = "coco" in args.data.dataset_name
+    # referred from https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/detect/val.py#L74
+    is_coco_dataset = (
+            isinstance(dataset_path, str)
+            and "coco" in dataset_path
+            and (args.data.val_set.endswith(f"{os.sep}val2017.txt") or args.data.val_set.endswith(f"{os.sep}test-dev2017.txt"))
+    )  # is COCO
     dataset = COCODataset(
         dataset_path=dataset_path,
         img_size=args.img_size,

diff --git a/train.py b/train.py
@@ -239,7 +239,12 @@ def train(args):
 
     # Create test function for run eval while train
     if args.run_eval:
-        is_coco_dataset = "coco" in args.data.dataset_name
+        # referred from https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/detect/val.py#L74
+        is_coco_dataset = (
+                isinstance(args.data.val_set, str)
+                and "coco" in args.data.val_set
+                and (args.data.val_set.endswith(f"{os.sep}val2017.txt") or args.data.val_set.endswith(f"{os.sep}test-dev2017.txt"))
+        )  # is COCO
         test_fn = partial(
             test,
             task=args.task,