Update main.py

main.py

@@ -16,6 +16,28 @@ PY_MODULES = {
     "wrinkle_unet.py": "WrinkleDetector"
 }
 
+def load_model(token):
+    repo_id = "IFMedTech/Skin-Analysis"
+    filename = "model/wrinkles_unet_v1.pth"
+    # token = os.environ.get("HUGGINGFACE_HUB_TOKEN")  # Set this env var with your token
+    
+    # if not token:
+    #     raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable is required for private repo access.")
+    
+    model_path = hf_hub_download(repo_id=repo_id, filename=filename, token=token)
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = smp.Unet(
+        encoder_name="resnet34",
+        encoder_weights=None,
+        in_channels=3,
+        classes=1
+    )
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.to(device)
+    model.eval()
+    return model, device
+
 def dynamic_import(module_path, class_name):
     spec = importlib.util.spec_from_file_location(class_name, module_path)
     module = importlib.util.module_from_spec(spec)
@@ -65,7 +87,8 @@ def analyze_skin(image: np.ndarray, analysis_type: str) -> np.ndarray:
         # else:
         #     print(f"Oiliness detection error: {result.get('error')}")
     elif analysis_type == "Wrinkles":
-        detector = detector_classes["WrinkleDetector"](image)
+        model, device = load_model(token)
+        detector = detector_classes["WrinkleDetector"](image, model, device)
         result = detector.predict_json()
         if result.get("detected") is not None:
             output = detector.draw_json(result)