Update app.py

app.py

@@ -8,16 +8,23 @@ from tensorflow.keras.models import load_model
 from torchvision.models import resnet18
 import os
 import requests
+import mediapipe as mp
 import cv2
 
-# Title
+# App title
 st.title("🧠 Stroke Patient Pain Intensity Detector")
 
-st.markdown("""
-Upload a full-face image of a stroke patient. The app detects the **affected facial side** using a stroke classification model, then uses the **unaffected side** to predict **pain intensity (PSPI score)**.
-""")
-
-# Download models
+# Instructions
+st.markdown(
+    """
+    Upload a full-face image of a stroke patient.
+    The app will detect the **affected facial side** using a stroke classification model,
+    and then use the **unaffected side** to predict **pain intensity** (PSPI score).
+    """
+)
+st.write("🔧 Initializing and downloading models...")
+
+# Download and load models
 @st.cache_resource
 def download_models():
     model_urls = {
@@ -26,134 +33,117 @@ def download_models():
     }
     for filename, url in model_urls.items():
         if not os.path.exists(filename):
+            st.write(f"📥 Downloading {filename}...")
             r = requests.get(url)
             with open(filename, "wb") as f:
                 f.write(r.content)
-    if not os.path.exists("haarcascade_frontalface_default.xml"):
-        r = requests.get("https://raw.githubusercontent.com/opencv/opencv/master/data/haarcascades/haarcascade_frontalface_default.xml")
-        with open("haarcascade_frontalface_default.xml", "wb") as f:
-            f.write(r.content)
+            st.success(f"✅ {filename} downloaded.")
+        else:
+            st.write(f"✔️ {filename} already exists.")
+
     stroke_model = load_model("cnn_stroke_model.keras")
     pain_model = resnet18(weights=None)
     pain_model.fc = nn.Linear(pain_model.fc.in_features, 1)
     pain_model.load_state_dict(torch.load("pain_model.pth", map_location=torch.device("cpu")))
     pain_model.eval()
+
     return stroke_model, pain_model
 
 stroke_model, pain_model = download_models()
 
-# Rotation using eyes
-def auto_rotate_face(image, box):
-    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    x, y, w, h = box
-    roi = img_cv[y:y+h, x:x+w]
-    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
-    eye_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_eye.xml")
-    eyes = eye_cascade.detectMultiScale(gray, 1.1, 5)
-    if len(eyes) >= 2:
-        eyes = sorted(eyes[:2], key=lambda e: e[0])
-        (x1, y1, w1, h1), (x2, y2, w2, h2) = eyes
-        center1 = (x1 + w1 // 2, y1 + h1 // 2)
-        center2 = (x2 + w2 // 2, y2 + h2 // 2)
-        dx, dy = center2[0] - center1[0], center2[1] - center1[1]
-        angle = np.degrees(np.arctan2(dy, dx))
-        return image.rotate(-angle, center=(x + w // 2, y + h // 2), expand=True)
-    return image
-
-# Image transform
+# Preprocessing for pain model
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-uploaded = st.file_uploader("📂 Upload face image", type=["jpg", "jpeg", "png"])
-if uploaded:
-    full_image = Image.open(uploaded).convert("RGB")
-    st.image(full_image, caption="Uploaded Full-Face Image", use_column_width=True)
-
-    # Detect face
-    cv_image = cv2.cvtColor(np.array(full_image), cv2.COLOR_RGB2BGR)
-    face_cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
-    faces = face_cascade.detectMultiScale(cv_image, 1.1, 5)
-
-    if len(faces) == 0:
-        st.error("❌ No face detected. Please upload a clear frontal face image.")
-        st.stop()
-
-    # Rotate and redetect
-    x, y, w, h = faces[0]
-    rotated = auto_rotate_face(full_image, (x, y, w, h))
-    rotated_cv = cv2.cvtColor(np.array(rotated), cv2.COLOR_RGB2BGR)
-    faces2 = face_cascade.detectMultiScale(rotated_cv, 1.1, 5)
-
-    if len(faces2) == 0:
-        st.warning("⚠️ Redetection failed after rotation, using original box.")
-        x, y, w, h = faces[0]
-    else:
-        x, y, w, h = faces2[0]
-
-    # Add padding
-    pad = int(0.2 * max(w, h))
-    x1, y1 = max(0, x - pad), max(0, y - pad)
-    x2, y2 = min(rotated.width, x + w + pad), min(rotated.height, y + h + pad)
-    face_crop = rotated.crop((x1, y1, x2, y2))
-
-    # Split face
-    fw, fh = face_crop.size
-    mid = fw // 2
-    patient_right = face_crop.crop((0, 0, mid, fh))  # viewer's left
-    patient_left  = face_crop.crop((mid, 0, fw, fh)) # viewer's right
-
-    # Stroke model input
-    _, H, W, C = stroke_model.input_shape
-    stroke_input = face_crop.resize((W, H))
-    stroke_array = np.array(stroke_input).astype("float32") / 255.0
-    stroke_array = np.expand_dims(stroke_array, axis=0)
-
-    # Predict affected side
-    st.write("🧠 Predicting affected side...")
-    stroke_pred = stroke_model.predict(stroke_array)
-    stroke_raw = stroke_pred[0][0]
-    affected = int(np.round(stroke_raw))  # 0 = left, 1 = right
-
-    # Assign faces
-    if affected == 0:
-        affected_side = "left"
-        unaffected_side = "right"
-        unaffected_face = patient_right
-    else:
-        affected_side = "right"
-        unaffected_side = "left"
-        unaffected_face = patient_left
-
-    # Predict pain score
-    st.write("📈 Predicting PSPI pain score...")
-    tensor = transform(unaffected_face).unsqueeze(0)
-    with torch.no_grad():
-        output = pain_model(tensor)
-        raw_score = output.item()
-        pspi_score = max(0.0, min(raw_score, 6.0))
-
-    # Output
-    st.subheader("🔍 Prediction Results")
-    st.image(unaffected_face, caption="Unaffected Side Used for Pain Detection", width=300)
-    st.write(f"**🧭 Affected Side (face POV):** `{affected_side}`")
-    st.write(f"**✅ Unaffected Side (face POV):** `{unaffected_side}`")
-    st.write(f"**🎯 Predicted PSPI Pain Score:** `{pspi_score:.3f}`")
-    st.write(f"**📈 Raw Pain Model Output:** `{raw_score:.3f}`")
-    st.write(f"**📊 Stroke Model Raw Output:** `{stroke_raw:.4f}`")
-
-    st.markdown("""
-    ---
-    ### ℹ️ Stroke Model Output
-    - Output ∈ [0, 1]
-    - Closer to `0` = Left side affected
-    - Closer to `1` = Right side affected
-
-    ### ℹ️ PSPI Score (0–6)
-    - 0: No pain  
-    - 1–2: Mild  
-    - 3–4: Moderate  
-    - 5–6: Severe
-    """)
+# MediaPipe Face Detection
+mp_face = mp.solutions.face_detection
+mp_draw = mp.solutions.drawing_utils
+
+# Upload UI
+uploaded_file = st.file_uploader("📂 Upload a full-face image", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    st.write("📷 Image uploaded. Detecting face...")
+    full_image = Image.open(uploaded_file).convert("RGB")
+    img_np = np.array(full_image)
+
+    with mp_face.FaceDetection(model_selection=1, min_detection_confidence=0.6) as detector:
+        results = detector.process(cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR))
+
+        if not results.detections:
+            st.error("❌ No face detected. Please upload a clear frontal face image.")
+            st.stop()
+
+        # Use first detection
+        detection = results.detections[0]
+        bboxC = detection.location_data.relative_bounding_box
+        ih, iw, _ = img_np.shape
+        x = int(bboxC.xmin * iw)
+        y = int(bboxC.ymin * ih)
+        w = int(bboxC.width * iw)
+        h = int(bboxC.height * ih)
+
+        face_crop = full_image.crop((x, y, x + w, y + h))
+        st.image(full_image, caption="Uploaded Full-Face Image", use_column_width=True)
+
+        # Split halves (face POV)
+        fw, fh = face_crop.size
+        fmid = fw // 2
+        patient_right = face_crop.crop((0, 0, fmid, fh))    # viewer's left
+        patient_left = face_crop.crop((fmid, 0, fw, fh))    # viewer's right
+
+        # Stroke prediction input
+        _, H, W, C = stroke_model.input_shape
+        stroke_input = face_crop.resize((W, H))
+        stroke_array = np.array(stroke_input).astype("float32") / 255.0
+        stroke_array = np.expand_dims(stroke_array, axis=0)
+
+        st.write("🧠 Predicting affected side of the face...")
+        stroke_pred = stroke_model.predict(stroke_array)
+        stroke_raw = stroke_pred[0][0]
+        affected = int(np.round(stroke_raw))  # 0 = left affected, 1 = right affected
+
+        if affected == 0:
+            affected_side = "left"
+            unaffected_side = "right"
+            unaffected_face = patient_right
+        else:
+            affected_side = "right"
+            unaffected_side = "left"
+            unaffected_face = patient_left
+
+        # Pain prediction
+        st.write("📈 Predicting PSPI pain score from unaffected side...")
+        input_tensor = transform(unaffected_face).unsqueeze(0)
+        with torch.no_grad():
+            output = pain_model(input_tensor)
+            raw_score = output.item()
+            pspi_score = max(0.0, min(raw_score, 6.0))
+
+        # Display results
+        st.subheader("🔍 Prediction Results")
+        st.image(unaffected_face, caption="Unaffected Side Used for Pain Detection", width=300)
+        st.write(f"**🧭 Affected Side (face POV):** `{affected_side}`")
+        st.write(f"**✅ Unaffected Side (face POV):** `{unaffected_side}`")
+        st.write(f"**🎯 Predicted PSPI Pain Score:** `{pspi_score:.3f}`")
+        st.write(f"**📈 Raw Pain Model Output:** `{raw_score:.3f}`")
+        st.write(f"**📊 Stroke Model Raw Output:** `{stroke_raw:.4f}`")
+
+        st.markdown(
+            """
+            ---
+            ### ℹ️ Stroke Model Output
+            - Output is between `0` and `1`
+            - Closer to `0` = Left side is affected
+            - Closer to `1` = Right side is affected
+            
+            ### ℹ️ PSPI Score Scale
+            - `0`: No pain  
+            - `1–2`: Mild pain  
+            - `3–4`: Moderate pain  
+            - `5–6`: Severe pain
+            """
+        )