fix yolov5 export, fix ci workflow for scripts (#62)

* fix yolov5 export

* fix yolov5 script ci

* fix yolov5 conversion script

.github/workflows/ci.yml

@@ -96,6 +96,8 @@ jobs:
           python scripts/predict.py --source tests/data/ --model_path tests/data/models/mmdet_retinanet/retinanet_r50_fpn_2x_coco_20200131-fdb43119.pth --config_path tests/data/models/mmdet_retinanet/retinanet_r50_fpn_1x_coco_v280.py
           python scripts/predict.py --source tests/data/coco_utils/terrain1.jpg --novisual --pickle --crop --model_path tests/data/models/mmdet_retinanet/retinanet_r50_fpn_2x_coco_20200131-fdb43119.pth --config_path tests/data/models/mmdet_retinanet/retinanet_r50_fpn_1x_coco_v280.py
           python scripts/predict.py --source tests/data/coco_utils/ --coco_file tests/data/coco_utils/combined_coco.json --model_path tests/data/models/mmdet_retinanet/retinanet_r50_fpn_2x_coco_20200131-fdb43119.pth --config_path tests/data/models/mmdet_retinanet/retinanet_r50_fpn_1x_coco_v280.py
+          # coco2yolov5
+          python scripts/coco2yolov5.py --source tests/data/coco_utils/ --coco_file tests/data/coco_utils/combined_coco.json --train_split 0.9
 
       - name: Test SAHI scripts for mmdet > 2.8.0
         if: matrix.mmdet-version == '2.11.0'

sahi/utils/coco.py

@@ -1028,11 +1028,17 @@ def split_coco_as_train_val(
         val_images = shuffled_images[num_train:]
 
         # form train val coco objects
-        train_coco = Coco(name=self.name if self.name else "split" + '_train')
+        train_coco = Coco(
+            name=self.name if self.name else "split" + '_train',
+            image_dir=self.image_dir
+        )
         train_coco.images = train_images
         train_coco.categories = self.categories
 
-        val_coco = Coco(name=self.name if self.name else "split" + '_val')
+        val_coco = Coco(
+            name=self.name if self.name else "split" + '_val',
+            image_dir=self.image_dir
+        )
         val_coco.images = val_images
         val_coco.categories = self.categories
 
@@ -1042,14 +1048,12 @@ def split_coco_as_train_val(
             "val_coco": val_coco,
         }
 
-    def export_as_yolov5(self, image_dir, output_dir, train_split_rate=1, numpy_seed=0):
+    def export_as_yolov5(self, output_dir, train_split_rate=1, numpy_seed=0):
         """
         Exports current COCO dataset in ultralytics/yolov5 format.
         Creates train val folders with image symlinks and txt files and a data yaml file.
 
         Args:
-            image_dir: str
-                Source image directory that contains coco images.
             output_dir: str
                 Export directory.
             train_split_rate: float
@@ -1099,11 +1103,11 @@ def export_as_yolov5(self, image_dir, output_dir, train_split_rate=1, numpy_seed
         # create image symlinks and annotation txts
         if split_mode in ["TRAINVAL", "TRAIN"]:
             export_yolov5_images_and_txts_from_coco_object(
-                image_dir, output_dir=train_dir, coco=train_coco
+                output_dir=train_dir, coco=train_coco
             )
         if split_mode in ["TRAINVAL", "VAL"]:
             export_yolov5_images_and_txts_from_coco_object(
-                image_dir, output_dir=val_dir, coco=val_coco
+                output_dir=val_dir, coco=val_coco
             )
 
         # create yolov5 data yaml
@@ -1136,14 +1140,12 @@ def get_subsampled_coco(self, subsample_ratio=10):
 
 
 def export_yolov5_images_and_txts_from_coco_object(
-    image_dir, output_dir, coco
+    output_dir, coco
 ):
     """
     Creates image symlinks and annotation txts in yolo format from coco dataset.
 
     Args:
-        image_dir: str
-            Source image directory that contains coco images.
         output_dir: str
             Export directory.
         coco: sahi.utils.coco.Coco
@@ -1152,16 +1154,21 @@ def export_yolov5_images_and_txts_from_coco_object(
 
     for image in tqdm(coco.images):
         # Create a symbolic link pointing to src named dst
-        src = os.path.abspath(os.path.join(image_dir, image.file_name))
-        dst = os.path.join(output_dir, image.file_name)
+        if Path(image.file_name).is_file():
+            src = os.path.abspath(image.file_name)
+        else:
+            assert coco.image_dir, "You have to specify image_dir " \
+                "of Coco object for yolov5 conversion."
+            src = os.path.abspath(str(Path(coco.image_dir) / image.file_name))
+        dst = str(Path(output_dir) / ((coco.name if coco.name else "") + "_" + Path(image.file_name).name))
         os.symlink(src, dst)
         # calculate annotation normalization ratios
         width = image.width
         height = image.height
         dw = 1.0 / (width)
         dh = 1.0 / (height)
         # set annotation filepath
-        file_name = get_base_filename(image.file_name)[1]
+        file_name = Path(image.file_name).stem
         yolo_annotation_filepath = "{}.txt".format(os.path.join(output_dir, file_name))
         # create annotation file
         annotations = image.annotations

scripts/coco2yolov5.py

@@ -33,10 +33,9 @@
     # increment run
     save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=False))
     # load coco dict
-    coco = Coco.from_coco_dict_or_path(coco_dict_or_path=opt.coco_file)
+    coco = Coco.from_coco_dict_or_path(coco_dict_or_path=opt.coco_file, image_dir=opt.source, mp=True)
     # export as yolov5
     coco.export_as_yolov5(
-        image_dir=opt.source,
         output_dir=str(save_dir),
         train_split_rate=opt.train_split,
         numpy_seed=opt.seed,

tests/test_cocoutils.py

@@ -311,9 +311,9 @@ def test_coco2yolo(self):
         output_dir = "tests/data/coco2yolo/"
         if os.path.isdir(output_dir):
             shutil.rmtree(output_dir)
-        coco = Coco.from_coco_dict_or_path(coco_dict_path)
+        coco = Coco.from_coco_dict_or_path(coco_dict_path, image_dir=image_dir)
         coco.export_as_yolov5(
-            image_dir, output_dir=output_dir, train_split_rate=0.5, numpy_seed=0
+            output_dir=output_dir, train_split_rate=0.5, numpy_seed=0
         )
 
     def test_update_categories(self):