From 7105f34f4421156d8d1f6added8f06c69f175eb7 Mon Sep 17 00:00:00 2001 From: SWHL Date: Tue, 23 Apr 2024 18:24:20 +0800 Subject: [PATCH] Support yolov8 to yolov5 --- label_convert/yolov8_to_yolov5.py | 65 +++++++++++++++++++++++-------- 1 file changed, 48 insertions(+), 17 deletions(-) diff --git a/label_convert/yolov8_to_yolov5.py b/label_convert/yolov8_to_yolov5.py index b1a5f65..f8e7a70 100644 --- a/label_convert/yolov8_to_yolov5.py +++ b/label_convert/yolov8_to_yolov5.py @@ -4,15 +4,21 @@ import argparse import shutil from pathlib import Path -from typing import Tuple, Union +from typing import List, Tuple, Union +import yaml from tqdm import tqdm ValueType = Union[str, Path, None] class YOLOV8ToYOLOV5: - def __init__(self, data_dir: ValueType = None, save_dir: ValueType = None): + def __init__( + self, + data_dir: ValueType = None, + save_dir: ValueType = None, + yaml_path: ValueType = None, + ): if data_dir is None: raise ValueError("data_dir must not be None") self.data_dir = Path(data_dir) @@ -25,10 +31,12 @@ def __init__(self, data_dir: ValueType = None, save_dir: ValueType = None): self.verify_exists(self.label_dir) if save_dir is None: - save_dir = self.data_dir.parent / f"{Path(self.data_dir).name}_yolov8" + save_dir = self.data_dir.parent / f"{Path(self.data_dir).name}_yolov5" self.save_dir = save_dir self.mkdir(self.save_dir) + self.yaml_path = yaml_path + self.save_img_dir = save_dir / "images" self.mkdir(self.save_img_dir) @@ -40,24 +48,29 @@ def __call__(self, mode_list: Tuple[str] = ("train", "val")): raise ValueError("mode_list is empty!!") for mode in tqdm(mode_list): - txt_path = self.data_dir / f"{mode}.txt" - self.verify_exists(txt_path) - img_list = self.read_txt(txt_path) - - save_mode_img_dir = self.save_img_dir / mode - self.mkdir(save_mode_img_dir) - - save_mode_label_dir = self.save_label_dir / mode - self.mkdir(save_mode_label_dir) + img_dir = self.img_dir / mode + img_list = list(img_dir.iterdir()) + img_relative_list = [] # copy images to new img dir for img_path in img_list: - img_full_path = self.data_dir / img_path - shutil.copy(img_full_path, save_mode_img_dir) + shutil.copy(img_path, self.save_img_dir) - label_path = self.label_dir / Path(img_path).with_suffix(".txt").name - shutil.copy(label_path, save_mode_label_dir) + label_path = self.label_dir / mode / img_path.with_suffix(".txt").name + shutil.copy(label_path, self.save_label_dir) + new_img_path = Path("images") / img_path.name + img_relative_list.append(new_img_path) + + txt_path = self.save_dir / f"{mode}.txt" + self.write_txt(txt_path, img_relative_list) + + class_txt_path = self.save_dir / "classes.txt" + class_content = "" + if self.yaml_path is not None: + yaml_data = self.read_yaml(self.yaml_path) + class_content = list(yaml_data["names"].values()) + self.write_txt(class_txt_path, class_content) print(f"Successfully convert, detail in {self.save_dir}") @staticmethod @@ -75,6 +88,24 @@ def verify_exists(file_path: Union[Path, str]): if not Path(file_path).exists(): raise FileNotFoundError(f"The {file_path} is not exists!!!") + @staticmethod + def write_txt(save_path: str, content: List[str], mode="w") -> None: + if not isinstance(save_path, str): + save_path = str(save_path) + + if isinstance(content, str): + content = [content] + + with open(save_path, mode, encoding="utf-8") as f: + for value in content: + f.write(f"{value}\n") + + @staticmethod + def read_yaml(yaml_path): + with open(yaml_path, "rb") as f: + data = yaml.load(f, Loader=yaml.Loader) + return data + def main(): parser = argparse.ArgumentParser("Datasets converter from YOLOV5 to COCO") @@ -91,7 +122,7 @@ def main(): parser.add_argument("--yaml_path", type=str, default=None) args = parser.parse_args() - converter = YOLOV8ToYOLOV5(args.data_dir, args.save_dir) + converter = YOLOV8ToYOLOV5(args.data_dir, args.save_dir, args.yaml_path) converter(mode_list=args.mode_list.split(","))