Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,162 +1,324 @@
 import streamlit as st
 import numpy as np
-import torch
 import cv2
-import os
-from PIL import Image
+import torch
 import tempfile
+from PIL import Image
 from tensorflow.keras.models import load_model
 from transformers import pipeline, AutoProcessor, LlavaForConditionalGeneration
-from huggingface_hub import login
+import io
+import os  # Import the os module to access environment variables
+from ultralytics import YOLO  # Import YOLO from ultralytics
 
-# Streamlit UI setup
-st.set_page_config(page_title="SmartHeal Wound Agent", layout="wide")
-st.title("🩹 SmartHeal: Agentic Wound Care Assistant")
+# --- Page Configuration (Best practice: call this first) ---
+st.set_page_config(
+    page_title="SmartHeal Wound Care Agent",
+    page_icon="🩹",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
 
-# ---- Load all models securely with HF token + local weights ----
+# --- Model Loading (Cached for performance) ---
 @st.cache_resource
 def load_all_models():
-    hf_token = os.environ.get("HUGGINGFACE_TOKEN")
-    if not hf_token:
-        raise ValueError("❌ HUGGINGFACE_TOKEN not found in environment")
-
-    login(token=hf_token)
-
-    # Detection model from local path
-    yolo_path = "/app/best.pt"
-    if not os.path.exists(yolo_path):
-        raise FileNotFoundError(f"Detection model not found: {yolo_path}")
-    detection_model = torch.hub.load("ultralytics/yolov5", "custom", path=yolo_path, force_reload=False)
-
-    # Segmentation model from local path
-    seg_path = "/app/segmentation model.h5"
-    if not os.path.exists(seg_path):
-        raise FileNotFoundError(f"Segmentation model not found: {seg_path}")
-    segmentation_model = load_model(seg_path, compile=False)
-
-    # Classification model from HuggingFace (private model allowed)
-    classification_pipe = pipeline(
-        "image-classification",
-        model="Hemg/Wound-classification",
-        use_auth_token=hf_token
-    )
-
-    # Med-Gemma (multimodal LLM)
-    model_id = "google/medgemma-4b-it"
-    processor = AutoProcessor.from_pretrained(model_id, token=hf_token)
-    med_model = LlavaForConditionalGeneration.from_pretrained(
-        model_id,
-        torch_dtype=torch.bfloat16,
-        token=hf_token,
-        low_cpu_mem_usage=True
-    ).to("cuda")
-
-    return detection_model, segmentation_model, classification_pipe, med_model, processor
-
-# ---- Inference Helper Functions ----
-def detect_wound(yolo_model, image_cv):
-    results = yolo_model(image_cv)
-    boxes = results.xyxy[0].cpu().numpy()
-    if len(boxes) == 0:
-        return None, None
-    x1, y1, x2, y2 = map(int, boxes[0][:4])
-    return image_cv[y1:y2, x1:x2], (x1, y1, x2, y2)
-
-def segment_wound(seg_model, region):
-    resized = cv2.resize(region, (256, 256)) / 255.0
-    pred = seg_model.predict(np.expand_dims(resized, axis=0))[0]
-    return (pred[:, :, 0] > 0.5).astype(np.uint8)
-
-def estimate_area(mask, px_per_cm=38):
-    return round(np.sum(mask > 0) / (px_per_cm ** 2), 2)
-
-def classify_wound(pipeline, region):
-    image = Image.fromarray(region)
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
-        image.save(tmp.name)
-        label = pipeline(tmp.name)[0]["label"]
-    os.unlink(tmp.name)
-    return label
-
-def generate_medgemma_response(image, processor, model, patient_info, area_cm2, wound_type):
-    messages = [
-        {
-            "role": "system",
-            "content": [{"type": "text", "text": "You are a wound care expert."}]
-        },
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": f"""Patient Info:
-- Age: {patient_info['age']}
-- Diabetic: {patient_info['diabetic']}
-- Wound Type: {wound_type}
-- Area: {area_cm2} cm²
-- Signs of infection: {patient_info['infection']}
-
-Please provide:
-1. Wound assessment
-2. Recommended treatment
-3. Cleaning & dressing method
-4. Red flags to monitor
-5. Follow-up schedule"""},
-                {"type": "image", "image": image}
-            ]
-        }
-    ]
-
-    input_ids = processor.apply_chat_template(messages, return_tensors="pt").to(model.device)
-    output = model.generate(input_ids, max_new_tokens=1000)
-    response = processor.decode(output[0], skip_special_tokens=True)
-
-    return response.split("ASSISTANT:")[-1].strip()
-
-# ---- Load models and run UI ----
-with st.spinner("🔄 Loading models..."):
-    yolo_model, seg_model, classify_pipe, med_model, processor = load_all_models()
-
-uploaded_file = st.file_uploader("📤 Upload a clear wound image", type=["jpg", "jpeg", "png"])
-with st.form("patient_form"):
-    age = st.number_input("Patient Age", min_value=1, max_value=120)
-    diabetic = st.radio("Diabetic?", ["Yes", "No"])
-    infection = st.radio("Visible infection?", ["Yes", "No"])
-    submit = st.form_submit_button("🚀 Analyze")
-
-if uploaded_file and submit:
-    image = Image.open(uploaded_file).convert("RGB")
-    image_cv = np.array(image)
-
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    with st.spinner("🧠 Analyzing image..."):
-        region, box = detect_wound(yolo_model, image_cv)
-        if region is None:
-            st.error("❌ No wound detected.")
+    """Loads all required models and pipelines into memory once."""
+    try:
+        # Get Hugging Face token from environment variable
+        hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
+        if not hf_token:
+            st.error("Fatal Error: Hugging Face token not found in environment variables (HF_TOKEN or HUGGING_FACE_HUB_TOKEN).")
             st.stop()
 
-        mask = segment_wound(seg_model, region)
-        area_cm2 = estimate_area(mask)
-        wound_type = classify_wound(classify_pipe, region)
-
-        response = generate_medgemma_response(
-            image=image,
-            processor=processor,
-            model=med_model,
-            patient_info={"age": age, "diabetic": diabetic, "infection": infection},
-            area_cm2=area_cm2,
-            wound_type=wound_type
+        # YOLOv8 detection model (using user's specified path)
+        detection_model = YOLO("/home/ubuntu/upload/best(1).pt")  # Load YOLOv8 model
+
+        # Segmentation model (using user's specified path)
+        segmentation_model = load_model("/home/ubuntu/upload/segmentation_model.h5", compile=False)
+
+        # Classification model (using user's specified model ID)
+        # Some pipelines might require token for private models or rate limits
+        classification_pipe = pipeline("image-classification", model="Hemg/Wound-classification", token=hf_token)
+
+        # Med-Gemma for analysis (using user's specified model ID)
+        medgemma_model_id = "google/medgemma-4b-it"
+        medgemma_processor = AutoProcessor.from_pretrained(medgemma_model_id, token=hf_token)
+        medgemma_model = LlavaForConditionalGeneration.from_pretrained(
+            medgemma_model_id,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            token=hf_token  # Pass the token here
         )
+        medgemma_model.to("cuda")  # Move model to GPU
+
+        return detection_model, segmentation_model, classification_pipe, medgemma_model, medgemma_processor
+    except Exception as e:
+        st.error(f"Fatal Error: Could not load models. Please check model paths, dependencies, and Hugging Face token. Details: {e}")
+        st.stop()
+
+
+# --- Agent Class ---
+class WoundCareAgent:
+    """An agentic class to encapsulate the wound care analysis pipeline."""
+    def __init__(self, models):
+        self.yolo_model, self.seg_model, self.classify_pipe, self.medgemma_model, self.medgemma_processor = models
+        self.px_per_cm = 38  # Example value, should be calibrated for real-world use
+
+    def detect_wound(self, image_cv):
+        """Detects the wound region using YOLOv8."""
+        st.session_state.messages.append({"role": "assistant", "content": "Detecting wound..."})
+        results = self.yolo_model(image_cv)  # Use YOLOv8 model directly
+        boxes = results[0].boxes.xyxy.cpu().numpy()  # Access boxes from YOLOv8 results
+        if len(boxes) == 0:
+            return None, None
+        # Assuming the largest bounding box is the wound (or the first detected)
+        box = boxes[0]
+        x1, y1, x2, y2 = map(int, box[:4])
+        detected_region = image_cv[y1:y2, x1:x2]
+        return detected_region, (x1, y1, x2, y2)
+
+    def segment_wound(self, detected_region):
+        """Segments the wound from the detected region using the provided segmentation model."""
+        st.session_state.messages.append({"role": "assistant", "content": "Segmenting wound area..."})
+        # Resize for segmentation model input
+        resized = cv2.resize(detected_region, (256, 256)) / 255.0
+        input_tensor = np.expand_dims(resized, axis=0)
+        pred_mask = self.seg_model.predict(input_tensor)[0]
+        binary_mask = (pred_mask[:, :, 0] > 0.5).astype(np.uint8)
+        return binary_mask
+
+    def estimate_area(self, mask):
+        """Estimates the area of the wound from the mask."""
+        st.session_state.messages.append({"role": "assistant", "content": "Estimating wound area..."})
+        pixel_area = np.sum(mask > 0)
+        area_cm2 = pixel_area / (self.px_per_cm ** 2)
+        return round(area_cm2, 2)
+
+    def classify_wound(self, detected_region):
+        """Classifies the type of the wound using the provided classification pipeline."""
+        st.session_state.messages.append({"role": "assistant", "content": "Classifying wound type..."})
+        try:
+            # Convert numpy array to PIL Image for the pipeline
+            pil_image = Image.fromarray(detected_region)
+            # Save to a temporary file for the pipeline, as it expects a file path or PIL Image
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
+                pil_image.save(tmp.name)
+                tmp_path = tmp.name
+            result = self.classify_pipe(tmp_path)
+            import os
+            os.unlink(tmp_path)  # Clean up temporary file
+            return result[0]["label"]
+        except Exception as e:
+            st.warning(f"Could not classify wound type: {e}")
+            return "Unknown"
+
+    def generate_recommendations(self, image, patient_info, analysis_results):
+        """Generates a detailed assessment and treatment plan using Med-Gemma."""
+        st.session_state.messages.append({"role": "assistant", "content": "Generating expert recommendations with Med-Gemma..."})
+
+        # Prepare the image for Med-Gemma (ensure it's a PIL Image)
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        # Construct the prompt for Med-Gemma
+        prompt_text = f"""Patient Info:
+- Age: {patient_info["age"]}
+- Diabetic: {patient_info["diabetic"]}
+- Wound Type: {analysis_results["wound_type"]}
+- Area: {analysis_results["area_cm2"]}
+- Signs of infection: {patient_info["infection"]}
+
+Please act as a highly experienced wound care specialist. Provide a comprehensive wound assessment and a detailed treatment plan. Structure your response clearly with the following sections:
+
+1.  **Wound Assessment:** Describe the wound's characteristics, potential causes, and current state based on the image and provided data.
+2.  **Recommended Treatment Plan:** Outline a primary course of action, including general principles of wound management.
+3.  **Cleaning and Dressing Protocol:** Provide specific, step-by-step instructions for wound cleaning and appropriate dressing choices.
+4.  **Red Flags & When to Seek Professional Medical Attention:** List critical signs and symptoms that indicate complications or require immediate consultation with a doctor or wound care nurse.
+5.  **Follow-up Schedule:** Suggest a realistic timeline for wound reassessment and monitoring progress.
+
+**Important Disclaimer:** This information is for educational purposes only and should not replace professional medical advice. Always consult a qualified healthcare provider for diagnosis and treatment.
+"""
+
+        # Med-Gemma expects messages in a specific format for multi-modal input
+        messages = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": "You are a wound care expert."}]
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt_text},
+                    {"type": "image", "image": image}
+                ]
+            }
+        ]
+
+        # Convert messages to input_ids using the processor's chat template
+        input_ids = self.medgemma_processor.apply_chat_template(messages, return_tensors="pt").to(self.medgemma_model.device)
+
+        # Generate response
+        output = self.medgemma_model.generate(input_ids, max_new_tokens=1000, do_sample=True, temperature=0.7)
+        response = self.medgemma_processor.decode(output[0], skip_special_tokens=True)
+
+        # Extract only the assistant's response part
+        if "ASSISTANT:" in response:
+            assistant_response = response.split("ASSISTANT:", 1)[1].strip()
+        else:
+            assistant_response = response.strip()
+
+        return assistant_response
+
+    def run_full_analysis(self, image, patient_info):
+        """Executes the entire analysis pipeline."""
+        st.session_state.messages.append({"role": "assistant", "content": "Starting analysis..."})
+        image_cv = np.array(image.convert("RGB"))
+
+        try:
+            detected_region, box = self.detect_wound(image_cv)
+            if detected_region is None:
+                st.error("Agent Error: No wound could be detected in the image. Please try another image.")
+                st.session_state.clear()
+                return None
+        except Exception as e:
+            st.error(f"Error during wound detection: {e}")
+            st.session_state.clear()
+            return None
 
-    st.success("✅ Analysis Complete")
-    st.markdown("### 📋 Med-Gemma Expert Recommendation")
-    st.info(response)
+        try:
+            mask_resized = self.segment_wound(detected_region)
+            area_cm2 = self.estimate_area(mask_resized)
+        except Exception as e:
+            st.error(f"Error during wound segmentation or area estimation: {e}")
+            st.session_state.clear()
+            return None
 
-    # Optional metrics & image
-    x1, y1, x2, y2 = box
-    overlay = image_cv.copy()
-    resized_mask = cv2.resize(mask, (x2 - x1, y2 - y1))
-    overlay[y1:y2, x1:x2][resized_mask > 0] = [255, 0, 0]
-    st.image(overlay, caption="Detection + Segmentation Overlay", use_column_width=True)
-    st.metric("Wound Area", f"{area_cm2} cm²")
-    st.metric("Wound Type", wound_type)
+        try:
+            wound_type = self.classify_wound(detected_region)
+        except Exception as e:
+            st.error(f"Error during wound classification: {e}")
+            st.session_state.clear()
+            return None
 
-    st.download_button("📥 Download Report", response, file_name="wound_report.txt")
+        analysis_results = {
+            "box": box,
+            "detected_region": detected_region,
+            "mask": mask_resized,
+            "area_cm2": area_cm2,
+            "wound_type": wound_type,
+        }
+
+        try:
+            recommendations = self.generate_recommendations(image, patient_info, analysis_results)
+            analysis_results["recommendations"] = recommendations
+        except Exception as e:
+            st.error(f"Error during recommendation generation: {e}")
+            st.session_state.clear()
+            return None
+
+        st.session_state.messages.append({"role": "assistant", "content": "Analysis complete. See results below."})
+        return analysis_results
+
+
+# --- UI Layout ---
+st.title("🩹 SmartHeal: The Agentic Wound Care Assistant")
+
+# Initialize session state
+if "analysis_results" not in st.session_state:
+    st.session_state.analysis_results = None
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# --- Sidebar for Inputs ---
+with st.sidebar:
+    st.header("📋 Patient & Image Input")
+    uploaded_file = st.file_uploader("1. Upload a clear wound image", type=["jpg", "jpeg", "png"])
+
+    with st.form("patient_form"):
+        st.write("2. Enter Patient Details")
+        age = st.number_input("Patient Age", min_value=1, max_value=120, value=50)
+        diabetic = st.radio("Is the patient diabetic?", ["No", "Yes"], index=0)
+        infection = st.radio("Are there visible signs of infection (e.g., pus, redness, swelling)?", ["No", "Yes"], index=0)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            submitted = st.form_submit_button("🚀 Analyze Wound", use_container_width=True)
+        with col2:
+            cleared = st.form_submit_button("❌ Clear", use_container_width=True)
+
+if cleared:
+    st.session_state.analysis_results = None
+    st.session_state.messages = []
+    st.rerun()
+
+# --- Main Content Area ---
+if submitted and uploaded_file:
+    # Load models and instantiate agent
+    models = load_all_models()
+    agent = WoundCareAgent(models)
+
+    # Store patient info
+    patient_info = {"age": age, "diabetic": diabetic, "infection": infection}
+
+    # Open image
+    image = Image.open(uploaded_file)
+
+    # Clear previous results and run new analysis
+    st.session_state.analysis_results = None
+    st.session_state.messages = []  # Clear messages for new analysis
+    st.session_state.messages.append({"role": "user", "content": "Analyzing the uploaded wound image..."})
+
+    with st.spinner("The SmartHeal Agent is at work..."):
+        st.session_state.analysis_results = agent.run_full_analysis(image, patient_info)
+
+# Display chat messages from the agent's process
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Display final results in tabs if analysis is complete
+if st.session_state.analysis_results:
+    results = st.session_state.analysis_results
+    image = Image.open(uploaded_file)  # Re-open image for display
+    image_cv = np.array(image.convert("RGB"))
+
+    st.header("✅ Analysis Complete")
+    tab1, tab2, tab3 = st.tabs(["📝 **Expert Recommendations**", "🔬 **Vision Analysis**", "📄 **Download Report**"])
+
+    with tab1:
+        st.markdown(results["recommendations"])
+
+    with tab2:
+        st.subheader("Wound Detection & Segmentation")
+
+        # Create overlay
+        x1, y1, x2, y2 = results["box"]
+        overlay_image = image_cv.copy()
+
+        # Ensure mask_resized is the correct size for the detected region
+        mask_for_overlay = cv2.resize(results["mask"], (x2 - x1, y2 - y1), interpolation=cv2.INTER_NEAREST)
+
+        # Create a colored mask for blending
+        colored_mask_region = np.zeros_like(overlay_image[y1:y2, x1:x2])
+        colored_mask_region[mask_for_overlay > 0] = [255, 0, 0]  # Red color for wound
+
+        # Blend the original detected region with the colored mask
+        overlay_image[y1:y2, x1:x2] = cv2.addWeighted(overlay_image[y1:y2, x1:x2], 0.7, colored_mask_region, 0.3, 0)
+
+        # Draw bounding box
+        cv2.rectangle(overlay_image, (x1, y1), (x2, y2), (0, 255, 0), 2)  # Green box
+
+        st.image(overlay_image, caption="Detected Wound with Segmentation Overlay", use_column_width=True)
+        st.metric(label="Estimated Wound Area", value=f"{results['area_cm2']} cm²")
+        st.metric(label="Classified Wound Type", value=f"{results['wound_type']}")
+
+    with tab3:
+        st.subheader("Download Full Report")
+        st.download_button(
+            label="📥 Download as Text File",
+            data=results["recommendations"],
+            file_name=f"wound_report_{uploaded_file.name.split('.')[0]}.txt",
+            mime="text/plain"
+        )
+else:
+    # Show information prompting the user to upload an image and fill in patient details.
+    st.info("Please upload an image and patient details in the sidebar to start.")