Update app.py

app.py

@@ -1,70 +1,102 @@
-import gradio as gr
-from PIL import Image
-import numpy as np
-import cv2
-
-from model_utils import get_model, predict
-
-# ---- Config ----
-CLASS_NAMES = ["background", "Normal_Conjuctiva", "conjunctiva_pallor"]
-WEIGHTS_PATH = r"skin_pallor_segment/Saved_model/mask_rcnn_conjunctiva.pth"
-
-# Optional: cache model so it loads once (recommended for Gradio)
-_MODEL = None
-
-def get_cached_model():
-    global _MODEL
-    if _MODEL is None:
-        _MODEL = get_model(num_classes=3, weights_path=WEIGHTS_PATH)
-    return _MODEL
-
-def segment_image(pil_img):
-    """
-    pil_img comes from gr.Image(type="pil") => already a PIL.Image (or None).
-    Returns a numpy RGB image for gr.Image output.
-    """
-    if pil_img is None:
-        return None
-
-    image = pil_img.convert("RGB")
-
-    model = get_cached_model()
-    results = predict(model, image, device="cpu", class_names=CLASS_NAMES)
-
-    # Overlay masks/contours on the original image
-    image_np = np.array(image)  # RGB uint8
-
-    for res in results:
-        mask = res["mask"]  # expected float/0..1
-        label = res.get("label", "")
-
-        colored_mask = (mask > 0.5).astype(np.uint8) * 255
-        contours, _ = cv2.findContours(colored_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
-
-        cv2.drawContours(image_np, contours, -1, (0, 255, 0), 2)
-
-        if len(contours) > 0 and len(contours[0]) > 0:
-            x, y = contours[0][0][0]
-            cv2.putText(
-                image_np,
-                str(label),
-                (int(x), int(y) - 10),
-                cv2.FONT_HERSHEY_SIMPLEX,
-                0.7,
-                (255, 0, 0),
-                2,
-            )
-
-    return image_np  # RGB numpy array works with gr.Image
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Conjunctiva Segmentation - Mask R-CNN")
-
-    with gr.Row():
-        inp = gr.Image(type="pil", label="Upload Image (Preview)")
-        out = gr.Image(type="numpy", label="Segmented Output")
-
-    submit = gr.Button("Submit")
-    submit.click(fn=segment_image, inputs=inp, outputs=out)  # button triggers inference [web:14][web:19]
-
-demo.launch()
+import gradio as gr
+from PIL import Image
+import numpy as np
+import cv2
+import os
+import torch
+
+from huggingface_hub import hf_hub_download
+from model_utils import get_model, predict
+
+# --- Config ---
+CLASS_NAMES = ["background", "Normal_Conjuctiva", "conjunctiva_pallor"]
+
+# Private repo + file in your HF model
+REPO_ID = "IFMedTech/Pallor_Mask_RCNN_Model"   # e.g. "IFMedTech/Eye-Anaemia-Model"
+FILENAME = "mask_rcnn_conjunctiva.pth"  # e.g. "models/mask_rcnn_conjunctiva.pth"
+
+def get_weights_path():
+    """
+    Download .pth from private HF repo using token stored in secrets.
+    """
+    token = os.environ.get("HUGGINGFACE_TOKEN")
+    if not token:
+        raise ValueError("Please set HUGGINGFACE_TOKEN in the Space secrets for private model access.")
+
+    model_path = hf_hub_download(
+        repo_id=REPO_ID,
+        filename=FILENAME,
+        token=token,
+    )
+    return model_path
+
+# Optional: cache model so it loads once (recommended for Gradio)
+_MODEL = None
+
+def get_cached_model():
+    global _MODEL
+    if _MODEL is None:
+        weights_path = get_weights_path()
+        # If your get_model expects weights_path, keep as is
+        _MODEL = get_model(num_classes=3, weights_path=weights_path)
+        # If instead you want to load via torch directly, you can:
+        # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        # model = get_model(num_classes=3)
+        # state_dict = torch.load(weights_path, map_location=device)
+        # model.load_state_dict(state_dict)
+        # model.to(device)
+        # model.eval()
+        # _MODEL = model
+    return _MODEL
+
+
+def segment_image(pil_img):
+    """
+    pil_img comes from gr.Image(type="pil") => already a PIL.Image (or None).
+    Returns a numpy RGB image for gr.Image output.
+    """
+    if pil_img is None:
+        return None
+
+    image = pil_img.convert("RGB")
+
+    model = get_cached_model()
+    results = predict(model, image, device="cpu", class_names=CLASS_NAMES)
+
+    # Overlay masks/contours on the original image
+    image_np = np.array(image)  # RGB uint8
+
+    for res in results:
+        mask = res["mask"]  # expected float/0..1
+        label = res.get("label", "")
+
+        colored_mask = (mask > 0.5).astype(np.uint8) * 255
+        contours, _ = cv2.findContours(colored_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+
+        cv2.drawContours(image_np, contours, -1, (0, 255, 0), 2)
+
+        if len(contours) > 0 and len(contours[0]) > 0:
+            x, y = contours[0][0][0]
+            cv2.putText(
+                image_np,
+                str(label),
+                (int(x), int(y) - 10),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.7,
+                (255, 0, 0),
+                2,
+            )
+
+    return image_np  # RGB numpy array works with gr.Image
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Conjunctiva Segmentation - Mask R-CNN")
+
+    with gr.Row():
+        inp = gr.Image(type="pil", label="Upload Image (Preview)")
+        out = gr.Image(type="numpy", label="Segmented Output")
+
+    submit = gr.Button("Submit")
+    submit.click(fn=segment_image, inputs=inp, outputs=out)  # button triggers inference [web:14][web:19]
+
+demo.launch()