Update app.py

app.py

@@ -9,6 +9,7 @@ import cv2
 import os
 import glob
 import time
+from huggingface_hub import hf_hub_download
 
 # Classes PASCAL VOC
 CLASSES = [
@@ -21,12 +22,21 @@ np.random.seed(42)
 COLORS = np.random.randint(50, 255, size=(len(CLASSES), 3), dtype=np.uint8)
 
 DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
-MODEL_PATH = "checkpoint_epoch_50.pth.tar"
+MODEL_REPO_ID = "nathbns/yolov1_from_scratch"
+MODEL_FILENAME = "checkpoint_epoch_50.pth.tar"
+
+# Charger le modèle depuis Hugging Face Hub
+print(f"Chargement du modèle depuis {MODEL_REPO_ID}...")
+try:
+    model_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME)
+    print(f"Modèle téléchargé depuis Hugging Face Hub: {model_path}")
+except Exception as e:
+    print(f"Erreur lors du téléchargement: {e}")
+    print("Tentative de chargement local...")
+    model_path = MODEL_FILENAME
 
-# Charger le modèle
-print(f"Chargement du modèle depuis {MODEL_PATH}...")
 model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE)
-checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
+checkpoint = torch.load(model_path, map_location=DEVICE)
 model.load_state_dict(checkpoint["state_dict"])
 model.eval()
 print(f"Modèle chargé avec succès!")