| from __future__ import annotations | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pathlib import Path | |
| from huggingface_hub import hf_hub_download | |
| from tqdm.auto import tqdm | |
| from ultralytics import YOLO | |
| BINGSU = "Bingsu/adetailer" | |
| GUON = "guon/hand-eyes" | |
| ULTRALYTICS = "Ultralytics/YOLOv8" | |
| model_info = [ | |
| (BINGSU, "deepfashion2_yolov8s-seg.pt"), | |
| (BINGSU, "face_yolov8m.pt"), | |
| (BINGSU, "face_yolov8n.pt"), | |
| (BINGSU, "face_yolov8n_v2.pt"), | |
| (BINGSU, "face_yolov8s.pt"), | |
| (BINGSU, "face_yolov9c.pt"), | |
| (BINGSU, "hand_yolov8n.pt"), | |
| (BINGSU, "hand_yolov8s.pt"), | |
| (BINGSU, "hand_yolov9c.pt"), | |
| (BINGSU, "person_yolov8m-seg.pt"), | |
| (BINGSU, "person_yolov8n-seg.pt"), | |
| (BINGSU, "person_yolov8s-seg.pt"), | |
| (GUON, "Eyes.pt"), | |
| (GUON, "PitEyeDetailer-v1-seg.pt"), | |
| (GUON, "PitHandDetailer-v1-seg.pt"), | |
| (GUON, "female_breast-v4.2.pt"), | |
| (GUON, "full_eyes_detect_v1.pt"), | |
| (GUON, "lips_v1.pt"), | |
| (ULTRALYTICS, "yolov8l.pt"), | |
| (ULTRALYTICS, "yolov8m.pt"), | |
| (ULTRALYTICS, "yolov8n.pt"), | |
| (ULTRALYTICS, "yolov8s.pt"), | |
| (ULTRALYTICS, "yolov8x.pt"), | |
| ] | |
| def download(model: str, repo: str) -> str: | |
| return hf_hub_download(repo, model, local_dir="downloads") | |
| def export(path: str) -> str: | |
| model = YOLO(path) | |
| return model.export(format="onnx", dynamic=True, simplify=True) | |
| def main(): | |
| save = Path("models") | |
| save.mkdir(exist_ok=True) | |
| pbar = tqdm(model_info) | |
| fs = [] | |
| with ThreadPoolExecutor() as executor: | |
| for repo, model in model_info: | |
| fu = executor.submit(download, model, repo) | |
| fu.add_done_callback(lambda _: pbar.update()) | |
| fs.append(fu) | |
| paths = [fu.result() for fu in fs] | |
| pbar.close() | |
| for path in paths: | |
| onnx = export(path) | |
| target_dir = save / Path(path).stem | |
| (target_dir / "1").mkdir(exist_ok=True, parents=True) | |
| (target_dir / "config.pbtxt").touch() | |
| target = target_dir / "1" / "model.onnx" | |
| Path(onnx).rename(target) | |
| if __name__ == "__main__": | |
| main() | |