From 1590d550d367863c06dc29d92ab773b2b816a090 Mon Sep 17 00:00:00 2001
From: unknown <2660417142@qq.com>
Date: Thu, 19 Dec 2024 10:31:26 +0800
Subject: [PATCH] fix is_coco_dataset judgment condition

---
 deploy/mslite_predict.py                  | 7 ++++++-
 examples/finetune_carparts_seg/predict.py | 7 ++++++-
 examples/finetune_visdrone/predict.py     | 7 ++++++-
 test.py                                   | 7 ++++++-
 train.py                                  | 7 ++++++-
 5 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/deploy/mslite_predict.py b/deploy/mslite_predict.py
index 0f59088b..8086f632 100644
--- 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
index 6ddc829e..a243358a 100644
--- 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
index ac20fe16..f684e5d4 100644
--- 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
index 525299a8..a5ac42cc 100644
--- 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
index 1faf1e37..a68dadf6 100644
--- 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,