kk

Acnes_model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:aa18e19b1c76c33a2edc8c5a820497fe9ff4e023be88cabae83d1ce9e289da47
-size 97947414
+oid sha256:c58a6af6fb622d916e4aa788d9f620d59df4a42209f655ce4a9db49686c05697
+size 97920000

app.py

@@ -12,8 +12,17 @@ MODEL_PATH = "Acnes_model.pth"
 MASK_OPACITY = 0.9
 DEVICE = torch.device("cpu")
 
-# ----------------- LOAD MODEL -----------------
-model = torch.load(MODEL_PATH, map_location=DEVICE)
+# ----------------- Rebuild Model & Load Weights -----------------
+# Recreate your model (same as in training)
+model = smp.Unet(
+    encoder_name="resnet34",
+    encoder_weights=None,
+    in_channels=3,
+    classes=1
+)
+
+# Load weights only (safe)
+model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
 model.to(DEVICE)
 model.eval()
 
@@ -66,7 +75,6 @@ def predict(image):
     predicted_mask = predict_mask(model, input_tensor)
     overlayed_image = overlay_mask(original_image, predicted_mask, color=(255, 0, 0), alpha=MASK_OPACITY)
 
-    # Save to disk for classification model (transformers pipeline accepts paths)
     temp_path = "/tmp/temp_image.jpg"
     cv2.imwrite(temp_path, cv2.cvtColor(original_image, cv2.COLOR_RGB2BGR))
 
@@ -85,7 +93,7 @@ demo = gr.Interface(
         gr.Image(label="Segmentation Overlay"),
         gr.Text(label="Acne Severity Prediction")
     ],
-    title="Acne Segmentation & Severity Classification",
+    title="🧼 Acne Segmentation & Severity Classification",
     description="Upload a facial image to detect acne regions and predict severity level using UNet and a pretrained classifier."
 )