diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aa9758db7..639350292 100644 --- a/.github/workflows/ci.yml +++ b/.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' diff --git a/sahi/utils/coco.py b/sahi/utils/coco.py index 4e102fb02..ee05b9eb5 100644 --- a/sahi/utils/coco.py +++ b/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,8 +1154,13 @@ 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 @@ -1161,7 +1168,7 @@ def export_yolov5_images_and_txts_from_coco_object( 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 diff --git a/scripts/coco2yolov5.py b/scripts/coco2yolov5.py index e3e6c30cf..16acd6474 100644 --- a/scripts/coco2yolov5.py +++ b/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, diff --git a/tests/test_cocoutils.py b/tests/test_cocoutils.py index f856af50c..8e5682fe5 100644 --- a/tests/test_cocoutils.py +++ b/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):