feat: model download

app.py

@@ -8,7 +8,7 @@ def get_face_type(img):
     pred_binary = get_pred_binary(img)
     result = [f"{label}: {bool(pred)}" for label, pred in zip(TARGET_LABELS, pred_binary)]
     face_type = int(''.join(map(str, pred_binary)), 2)
-    result = f"face_type: {face_type}\n{"\n".join(result)}"
+    result = f"face_type: {face_type}\n{'\n'.join(result)}"
     return result
 
 

main.py

@@ -5,11 +5,12 @@ import torch
 from deepface import DeepFace
 from fastapi import FastAPI, HTTPException
 
-from model import MultiLabelClassifier
+from model import MultiLabelClassifier, ensure_model_downloaded
 
 app = FastAPI()
 
-model_path = "classifier.pth"
+model_path = "data/classifier.pth"
+ensure_model_downloaded(model_path)
 model = MultiLabelClassifier(embedding_dim=4096, hidden_dim=1024)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.load_state_dict(torch.load(model_path, weights_only=True))

model.py

@@ -1,10 +1,12 @@
 import logging
+import os
 import pickle
 from concurrent.futures import ProcessPoolExecutor, as_completed
 from pathlib import Path
 
 import pandas
 import pandas as pd
+import requests
 import torch
 from deepface import DeepFace
 from sklearn.metrics import accuracy_score, recall_score, f1_score
@@ -35,6 +37,19 @@ def load_df(target_labels: list[str]):
     return train_df, test_df
 
 
+def ensure_model_downloaded(model_path: str):
+    os.makedirs(os.path.dirname(model_path), exist_ok=True)
+    if not os.path.exists(model_path):
+        logging.warning("Model not found. Downloading from GitHub...")
+        response = requests.get("https://github.com/wyyadd/facetype/releases/download/1.0.0/classifier.pth")
+        if response.status_code != 200:
+            logging.error("Failed to download classifier.pth")
+            raise RuntimeError("Failed to download model.")
+        with open(model_path, "wb") as f:
+            f.write(response.content)
+        logging.info("Download complete.")
+
+
 class EmbeddingDataset(Dataset):
     def __init__(self, df: pandas.DataFrame, target_labels: list[str]):
         self.df = df