feat: model download and convert
Browse files
models.py
ADDED
@@ -0,0 +1,59 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from __future__ import annotations
|
2 |
+
|
3 |
+
from pathlib import Path
|
4 |
+
|
5 |
+
from huggingface_hub import hf_hub_download
|
6 |
+
from ultralytics import YOLO
|
7 |
+
from tqdm.auto import tqdm
|
8 |
+
|
9 |
+
REPOBINGSU = "Bingsu/adetailer"
|
10 |
+
REPOGUON = "guon/hand-eyes"
|
11 |
+
|
12 |
+
model_info = [
|
13 |
+
(REPOBINGSU, "deepfashion2_yolov8s-seg.pt"),
|
14 |
+
(REPOBINGSU, "face_yolov8m.pt"),
|
15 |
+
(REPOBINGSU, "face_yolov8n.pt"),
|
16 |
+
(REPOBINGSU, "face_yolov8n_v2.pt"),
|
17 |
+
(REPOBINGSU, "face_yolov8s.pt"),
|
18 |
+
(REPOBINGSU, "face_yolov9c.pt"),
|
19 |
+
(REPOBINGSU, "hand_yolov8n.pt"),
|
20 |
+
(REPOBINGSU, "hand_yolov8s.pt"),
|
21 |
+
(REPOBINGSU, "hand_yolov9c.pt"),
|
22 |
+
(REPOBINGSU, "person_yolov8m-seg.pt"),
|
23 |
+
(REPOBINGSU, "person_yolov8n-seg.pt"),
|
24 |
+
(REPOBINGSU, "person_yolov8s-seg.pt"),
|
25 |
+
(REPOGUON, "Eyes.pt"),
|
26 |
+
(REPOGUON, "PitEyeDetailer-v1-seg.pt"),
|
27 |
+
(REPOGUON, "PitHandDetailer-v1-seg.pt"),
|
28 |
+
(REPOGUON, "female_breast-v4.2.pt"),
|
29 |
+
(REPOGUON, "full_eyes_detect_v1.pt"),
|
30 |
+
(REPOGUON, "lips_v1.pt"),
|
31 |
+
]
|
32 |
+
|
33 |
+
|
34 |
+
def download(model: str, repo: str) -> str:
|
35 |
+
return hf_hub_download(
|
36 |
+
repo, model, local_dir="downloads", local_dir_use_symlinks=False
|
37 |
+
)
|
38 |
+
|
39 |
+
|
40 |
+
def export(path: str) -> str:
|
41 |
+
model = YOLO(path)
|
42 |
+
return model.export(format="onnx", dynamic=True, simplify=True)
|
43 |
+
|
44 |
+
|
45 |
+
def main():
|
46 |
+
save = Path("models")
|
47 |
+
save.mkdir(exist_ok=True)
|
48 |
+
for repo, model in tqdm(model_info):
|
49 |
+
path = download(model, repo)
|
50 |
+
onnx = export(path)
|
51 |
+
target_dir = save / Path(model).stem
|
52 |
+
(target_dir / "1").mkdir(exist_ok=True, parents=True)
|
53 |
+
(target_dir / "config.pbtxt").touch()
|
54 |
+
target = target_dir / "1" / "model.onnx"
|
55 |
+
Path(onnx).rename(target)
|
56 |
+
|
57 |
+
|
58 |
+
if __name__ == "__main__":
|
59 |
+
main()
|