File size: 2,066 Bytes
18b55d3 189300e 18b55d3 189300e 18b55d3 189300e 18b55d3 189300e 18b55d3 24c40e1 18b55d3 189300e 24c40e1 189300e 24c40e1 189300e 24c40e1 18b55d3 4c7fc08 18b55d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
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()
|