Referencing local models

app.py

@@ -66,24 +66,18 @@ def calibration(source_img, target_img):
         return corrected_img.astype(np.uint8)
 
 
-device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
 # Load YOLO detector
 #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_path = hf_hub_download(repo_id="sanpdy/yolo-flake-detector", filename="uark_detector_v3.pt", token=False)
-yolo = YOLO(yolo_path)
-yolo.conf = 0.5
-
-# Load classifier weights
-classifier_path = hf_hub_download(
-    repo_id="sanpdy/flake-classifier",
-    filename="flake_classifier_5layer.pth",
-    token=False
-)
-ckpt = torch.load(classifier_path, map_location=device)
+
+yolo = YOLO("models/uark_detector_v3.pt")
+torch.load("models/flake_classifier_5layer.pth")
+ckpt_path = "models/flake_classifier_5layer.pth"
+ckpt = torch.load(ckpt_path, map_location=device)
 
 num_classes   = len(ckpt["class_to_idx"])
 classifier    = FlakeLayerClassifier(

models/flake_classifier_5layer.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e5d8873b2708e53cdd0b603918728c75c86c4a219e92b4895d5cfc5c39cc275
+size 47208000

models/uark_detector_v3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccfd0d79fc4e721474669054bb0eb4d2f6fb6e19f34e60a74d597947d9f66163
+size 6132509