Update app.py

app.py

@@ -1,155 +1,98 @@
 import streamlit as st
 import torch
 import torch.nn as nn
-from torchvision.models import resnet18
 from torchvision import transforms
 from PIL import Image
 import numpy as np
 from tensorflow.keras.models import load_model
 import os
 import requests
-import cv2
 
-st.title("🧠 Stroke Patient Pain Intensity Detector")
-st.markdown("Upload a full-face image. The system will detect the affected side and use the other side to predict pain intensity.")
+st.title("Stroke Patient Pain Intensity Detector")
+st.markdown("Upload a full-face image of a stroke patient. The app will detect the affected side and predict pain intensity using the unaffected side.")
+st.write("🔧 App started. Preparing to download models...")
 
 @st.cache_resource
 def download_models():
     model_urls = {
         "cnn_stroke_model.keras": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/cnn_stroke_model.keras",
-        "left_side_pain_classifier.pth": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/left_side_pain_classifier.pth",
-        "right_side_pain_classifier.pth": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/right_side_pain_classifier.pth",
+        "pain_model.pth": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/pain_model.pth"
     }
 
     for filename, url in model_urls.items():
-        if not os.path.exists(filename) or os.path.getsize(filename) < 10000:
+        if not os.path.exists(filename):
             st.write(f"📥 Downloading {filename}...")
-            try:
-                r = requests.get(url)
-                r.raise_for_status()
-                with open(filename, "wb") as f:
-                    f.write(r.content)
-                st.success(f"✅ {filename} downloaded.")
-            except Exception as e:
-                st.error(f"❌ Could not download {filename}: {e}")
-                st.stop()
+            r = requests.get(url)
+            with open(filename, "wb") as f:
+                f.write(r.content)
+            st.write(f"✅ {filename} downloaded.")
         else:
             st.write(f"✔️ {filename} already exists.")
 
-    # Haar cascade
-    haar_path = "haarcascade_frontalface_default.xml"
-    haar_url = "https://raw.githubusercontent.com/opencv/opencv/master/data/haarcascades/haarcascade_frontalface_default.xml"
-    if not os.path.exists(haar_path):
-        r = requests.get(haar_url)
-        with open(haar_path, "wb") as f:
-            f.write(r.content)
-
-    # Load stroke model
-    try:
-        stroke_model = load_model("cnn_stroke_model.keras")
-    except Exception as e:
-        st.error(f"❌ Failed to load stroke model: {e}")
-        st.stop()
-
-    def load_pain_model(path):
-        model = resnet18(weights=None)
-        model.fc = nn.Linear(model.fc.in_features, 1)
-        model.load_state_dict(torch.load(path, map_location="cpu"))
-        model.eval()
-        return model
-
-    try:
-        left_model = load_pain_model("left_side_pain_classifier.pth")
-        right_model = load_pain_model("right_side_pain_classifier.pth")
-    except Exception as e:
-        st.error(f"❌ Error loading PyTorch pain models: {e}")
-        st.stop()
-
-    return stroke_model, left_model, right_model
-
-stroke_model, left_model, right_model = download_models()
+    st.write("📦 Loading models...")
+    stroke_model = load_model("cnn_stroke_model.keras")
+
+    class PainRegressor(nn.Module):
+        def __init__(self):
+            super(PainRegressor, self).__init__()
+            from torchvision.models import resnet18, ResNet18_Weights
+            self.base = resnet18(weights=ResNet18_Weights.DEFAULT)
+            num_features = self.base.fc.in_features
+            self.base.fc = nn.Linear(num_features, 1)
+        def forward(self, x):
+            return self.base(x)
+
+    pain_model = PainRegressor()
+    pain_model.load_state_dict(torch.load("pain_model.pth", map_location=torch.device('cpu')))
+    pain_model.eval()
+    st.write("✅ Models loaded.")
+
+    return stroke_model, pain_model
+
+stroke_model, pain_model = download_models()
 
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225])
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
 ])
 
-uploaded_file = st.file_uploader("📂 Upload a full-face image", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
+    st.write("📷 Image uploaded. Processing...")
     full_image = Image.open(uploaded_file).convert("RGB")
-    st.image(full_image, caption="📸 Uploaded Full-Face Image", use_column_width=True)
-
-    np_img = np.array(full_image)
-    gray = cv2.cvtColor(np_img, cv2.COLOR_RGB2GRAY)
-    face_cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
-    faces = face_cascade.detectMultiScale(gray, 1.1, 5)
-
-    if len(faces) == 0:
-        st.error("❌ No face detected. Please upload a clearer image.")
-        st.stop()
-
-    # Use first detected face
-    (x, y, w, h) = faces[0]
-    center_x = x + w // 2
-    half_width = w // 2
-
-    left_half_box = (
-        max(center_x - half_width, 0),
-        y,
-        center_x,
-        y + h
-    )
-    right_half_box = (
-        center_x,
-        y,
-        min(center_x + half_width, full_image.width),
-        y + h
-    )
-
-    left_half = full_image.crop(left_half_box)
-    right_half = full_image.crop(right_half_box)
-
-    stroke_input = full_image.resize((224, 224))
-    stroke_array = np.expand_dims(np.array(stroke_input) / 255.0, axis=0)
-
-    st.write("🧠 Predicting affected side...")
-    stroke_pred = stroke_model.predict(stroke_array)
-    affected = int(np.round(stroke_pred[0][0]))
+    st.image(full_image, caption="Uploaded Full-Face Image", use_column_width=True)
 
-    # Flip logic to match face's own point of view
-    unaffected_face = left_half if affected == 1 else right_half
-    selected_model = left_model if affected == 1 else right_model
+    w, h = full_image.size
+    mid = w // 2
+    left_face = full_image.crop((0, 0, mid, h))
+    right_face = full_image.crop((mid, 0, w, h))
 
-    st.write("📈 Predicting pain...")
-    with torch.no_grad():
-        tensor = transform(unaffected_face).unsqueeze(0)
-        output = selected_model(tensor)
-        prob = torch.sigmoid(output).item()
-        label = 1 if prob > 0.5 else 0
-
-    st.subheader("🔍 Prediction Result")
-    st.image(unaffected_face, caption="🧍 Unaffected Side Used", width=300)
-
-  # ✅ Fix: if 1 means left is affected, use right side for pain
-if affected == 1:
-    affected_side = "left"
-    unaffected_side = "right"
-    unaffected_face = right_half
-    selected_model = right_model
-else:
-    affected_side = "right"
-    unaffected_side = "left"
-    unaffected_face = left_half
-    selected_model = left_model
+    # 🔍 Automatically resize image based on model input
+    _, H, W, C = stroke_model.input_shape
+    st.write(f"🔍 Resizing uploaded image to: ({H}, {W}) for stroke model")
 
-    st.write(f"**Affected Side (Face POV): **{face_affected}**")
-    st.write(f"**Unaffected Side (Face POV): **{face_unaffected}**")
-    st.write(f"**Predicted PSPI Pain Score:** {prob:.3f}")
-    st.write(f"**Stroke model raw output:** {stroke_pred[0][0]}")
+    stroke_input = full_image.resize((W, H))
+    stroke_array = np.array(stroke_input).astype("float32") / 255.0
+    stroke_array = np.expand_dims(stroke_array, axis=0)
 
+    st.write("🧠 Running stroke model prediction...")
+    stroke_pred = stroke_model.predict(stroke_array)
+    affected = int(np.round(stroke_pred[0][0]))
 
 
+    unaffected_face = right_face if affected == 0 else left_face
+    unaffected_tensor = transform(unaffected_face).unsqueeze(0)
 
+    st.write("📈 Predicting pain score...")
+    with torch.no_grad():
+        output = pain_model(unaffected_tensor)
+        pspi_score = output.item()
+
+    st.subheader("Prediction Results")
+    st.image(unaffected_face, caption="Unaffected Side Used for Pain Detection", width=300)
+    st.write(f"**Affected side:** {'left' if affected == 0 else 'right'}")
+    st.write(f"**Unaffected side:** {'right' if affected == 0 else 'left'}")
+    st.write(f"**Predicted PSPI Pain Score:** {round(pspi_score, 3)}")