Using huggingface-hosted models

app.py

@@ -9,6 +9,8 @@ import gradio as gr
 from ultralytics import YOLO
 from transformers import ResNetModel
 import cv2
+from huggingface_hub import hf_hub_download
+
 
 class FlakeLayerClassifier(nn.Module):
     def __init__(self, num_materials, material_dim, num_classes=4, dropout_prob=0.1, freeze_cnn=False):
@@ -71,14 +73,18 @@ print(f"Using device: {device}")
 #yolo = YOLO("/home/sankalp/flake_classification/models/best.pt")
 #yolo = YOLO("/home/sankalp/yolo_flake_detection/yolo11n_synthetic_runs/exp1/weights/best.pt")
 #yolo = YOLO("/home/sankalp/yolo_flake_detection/yolo_runs/yolo11l_flake_runs/weights/best.pt")
-yolo = YOLO("models/yolo-flake-detector.pt")
+yolo_path = hf_hub_download(repo_id="sanpdy/yolo-flake-detector", filename="yolo-flake-detector-MSU.pt", token=False)
+yolo = YOLO(yolo_path)
 yolo.conf = 0.5
 
 # Load classifier weights
-ckpt = torch.load(
-    "/home/sankalp/flake_classification/models/flake_classifier.pth",
-    map_location=device
+classifier_path = hf_hub_download(
+    repo_id="sanpdy/flake-classifier",
+    filename="flake-classifier.pth",
+    token=False
 )
+ckpt = torch.load(classifier_path, map_location=device)
+
 num_classes   = len(ckpt["class_to_idx"])
 classifier    = FlakeLayerClassifier(
     num_materials=num_classes,