deploy

.streamlit/config.toml

@@ -0,0 +1,22 @@
+[theme]
+primaryColor = "#1f77b4"
+backgroundColor = "#ffffff"
+secondaryBackgroundColor = "#f0f2f6"
+textColor = "#262730"
+font = "sans serif"
+
+[client]
+showErrorDetails = false
+
+[logger]
+level = "info"
+
+[server]
+port = 7860
+address = "0.0.0.0"
+headless = true
+runOnSave = true
+maxUploadSize = 500
+
+[browser]
+gatherUsageStats = false

Dockerfile

@@ -1,20 +1,30 @@
-FROM python:3.13.5-slim
+FROM python:3.11-slim
 
 WORKDIR /app
 
+# Install system dependencies
 RUN apt-get update && apt-get install -y \
-    build-essential \
+    libsm6 libxext6 libxrender-dev \
     curl \
-    git \
     && rm -rf /var/lib/apt/lists/*
 
-COPY requirements.txt ./
-COPY src/ ./src/
+# Copy requirements
+COPY requirements-prod.txt .
 
-RUN pip3 install -r requirements.txt
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements-prod.txt
 
+# Copy application
+COPY . .
+
+# Create outputs directory for models
+RUN mkdir -p outputs
+
+# Expose Streamlit port
 EXPOSE 8501
 
-HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
+# Health check
+HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health || exit 1
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+# Run Streamlit
+CMD ["streamlit", "run", "streamlit_app.py", "--logger.level=info", "--server.port=7860", "--server.address=0.0.0.0"]

apps/patient_chat_app_cloud.py

@@ -0,0 +1,666 @@
+"""
+Streamlit-based Patient Chat Application for Fracture Detection and Diagnosis.
+CLOUD VERSION - Uses Hugging Face Inference API instead of Ollama
+
+Supports:
+1. Running individual agents (Diagnostic, Educational, Explainability, Knowledge)
+2. Running the complete workflow
+3. LLM-based Q&A via Hugging Face Inference API for patient education
+"""
+
+import os
+import sys
+import streamlit as st
+import requests
+import json
+import numpy as np
+from typing import Dict, Any, List
+from pathlib import Path
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+# --- Import the Agents ---
+from src.agents.diagnostic_agent import DiagnosticAgent
+from src.agents.educational_agent import EducationalAgent
+from src.agents.explain_agent import ExplainabilityAgent, generate_random_heatmap, calculate_heatmap_centroid
+from src.agents.knowledge_agent import KnowledgeAgent, MEDICAL_KNOWLEDGE_BASE
+from src.agents.cross_validation_agent import ModelEnsembleAgent
+from src.utils import get_device
+
+# --- Hugging Face Inference API Configuration ---
+# Try both uppercase and lowercase key names for flexibility
+HF_API_KEY = st.secrets.get("HUGGINGFACE_API_KEY", st.secrets.get("huggingface_api_key", ""))
+# HF_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
+HF_API_URL = "https://router.huggingface.co/v1/chat/completions"
+HF_HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
+
+# --- Constants ---
+CLASS_NAMES = ["Comminuted", "Greenstick", "Healthy", "Oblique", 
+               "Oblique Displaced", "Spiral", "Transverse", "Transverse Displaced"]
+NUM_CLASSES = len(CLASS_NAMES)
+IMG_SIZE = 224
+
+# --- Page Configuration ---
+st.set_page_config(
+    page_title="🦴 Fracture Detection AI System",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# --- Custom CSS for Better UI ---
+st.markdown("""
+<style>
+    .stTabs [data-baseweb="tab-list"] button {
+        font-size: 16px;
+        font-weight: bold;
+    }
+    .section-header {
+        font-size: 20px;
+        font-weight: bold;
+        margin-top: 20px;
+        margin-bottom: 10px;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+
+# ============================================================================
+# --- 1. Hugging Face Inference API Patient Interaction Agent ---
+# ============================================================================
+
+class PatientInteractionAgent:
+    """
+    Handles communication with Mistral 7B model via Hugging Face Inference API.
+    Free tier available with rate limiting.
+    """
+    def __init__(self, medical_summary: Dict[str, Any], patient_history: Dict[str, Any]):
+        """Initialize the agent with medical context."""
+        # --- Configuration Check ---
+        if not HF_API_KEY:
+            raise ValueError(
+                "⚠️ HUGGINGFACE_API_KEY not found in Streamlit Secrets. "
+                "Please set your Hugging Face API token in Settings > Secrets."
+            )
+
+        self.medical_summary = medical_summary
+        self.patient_history = patient_history
+        self.system_prompt = self._build_system_prompt()
+
+    def _build_system_prompt(self) -> str:
+        """Creates a detailed instruction set for the LLM (RAG Context)."""
+        guidelines = "\n- ".join(self.medical_summary.get('Guidelines', ["No specific guidelines available."]))
+
+        return f"""You are a highly compassionate, clear, and professional medical assistant. Your goal is to answer patient questions in natural language based ONLY on the following diagnostic information and patient history.
+
+RULES:
+1. Maintain a reassuring, non-technical, and empathetic tone suitable for a patient.
+2. Keep answers concise and address the patient's underlying concern.
+3. ALWAYS conclude your answer by advising the patient to consult their orthopedic specialist or doctor.
+
+--- DIAGNOSTIC INFORMATION ---
+Diagnosis: {self.medical_summary.get('Diagnosis')} (Confidence: {self.medical_summary.get('Ensemble_Confidence')})
+Definition: {self.medical_summary.get('Type')}
+Severity: {self.medical_summary.get('Severity')}
+Treatment Guidelines:
+{guidelines}
+
+--- PATIENT HISTORY ---
+Age: {self.patient_history.get('age')}
+Gender: {self.patient_history.get('gender')}
+Medical History: {self.patient_history.get('history')}"""
+
+    def get_response(self, query: str) -> str:
+        """Queries the Hugging Face Inference API with the patient's question."""
+        try:
+            # Format prompt for Mistral using [INST] tags
+            full_prompt = f"{self.system_prompt}\n\nPATIENT QUERY: {query}"
+            
+            # payload = {
+            #     "inputs": f"[INST] {full_prompt} [/INST]",
+            #     "parameters": {
+            #         "max_new_tokens": 512,
+            #         "return_full_text": False,
+            #         "temperature": 0.7,
+            #     }
+            # }
+            payload = {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": f"[INST] {full_prompt} [/INST]"
+                    }
+                ],
+                "model": "meta-llama/Llama-3.1-8B-Instruct:cerebras"
+            }
+
+            response = requests.post(
+                HF_API_URL,
+                headers=HF_HEADERS,
+                json=payload,
+                timeout=60
+            )
+            response.raise_for_status()
+
+            result = response.json()
+            result = result["choices"][0]["message"]
+
+            # Handle different response formats
+            if isinstance(result, list) and len(result) > 0:
+                return result[0].get("generated_text", "Error: Unexpected API response format.")
+            elif isinstance(result, dict) and "generated_text" in result:
+                return result["generated_text"]
+            elif isinstance(result, dict) and "content" in result:
+                return result["content"]
+            elif "error" in result:
+                # Handle API errors (e.g., model loading, rate limiting)
+                error_msg = result.get("error", "Unknown error")
+                if "rate_limit" in str(error_msg).lower():
+                    return "⚠️ API rate limit reached. Please wait a moment and try again."
+                return f"⚠️ API Error: {error_msg}"
+            else:
+                return "Error: Unknown API response format."
+
+        except requests.exceptions.Timeout:
+            return "⏱️ Request timed out. The model may be loading. Please try again."
+        except requests.exceptions.RequestException as e:
+            return f"❌ Network error: {str(e)}"
+        except Exception as e:
+            return f"❌ Unexpected error: {str(e)}"
+
+
+# ============================================================================
+# --- 2. Helper Functions ---
+# ============================================================================
+
+def save_uploaded_file(uploaded_file) -> str:
+    """Save uploaded file to a temporary location."""
+    if uploaded_file is None:
+        return None
+    
+    try:
+        import tempfile
+        # Create a temporary file in temp_uploads directory
+        temp_dir = Path("./temp_uploads")
+        temp_dir.mkdir(exist_ok=True)
+        
+        # Create temp file with proper extension
+        suffix = Path(uploaded_file.name).suffix or '.jpg'
+        with tempfile.NamedTemporaryFile(
+            dir=str(temp_dir),
+            suffix=suffix,
+            delete=False
+        ) as tmp_file:
+            tmp_file.write(uploaded_file.getbuffer())
+            return tmp_file.name  # Returns full path
+    except Exception as e:
+        st.error(f"Error saving file: {e}")
+        return None
+
+
+# ============================================================================
+# --- 3. Workflow Functions ---
+# ============================================================================
+
+def run_diagnostic_agent(image_path: str) -> Dict[str, Any]:
+    """Run the diagnostic agent on an image."""
+    try:
+        # Placeholder checkpoint path - in production, use actual model checkpoint
+        checkpoint_path = "./outputs/best_swin.pth"
+        
+        if not os.path.exists(checkpoint_path):
+            return {"error": f"Checkpoint not found at {checkpoint_path}"}
+        
+        agent = DiagnosticAgent(
+            checkpoint_path=checkpoint_path,
+            model_name='swin',
+            num_classes=NUM_CLASSES,
+            img_size=IMG_SIZE,
+            class_names=CLASS_NAMES
+        )
+        
+        result = agent.run_diagnosis(image_path)
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def run_ensemble_agent(image_path: str) -> Dict[str, Any]:
+    """Run the ensemble agent on an image."""
+    try:
+        checkpoints_dir = "./outputs"
+        
+        if not os.path.exists(checkpoints_dir):
+            return {"error": f"Checkpoints directory not found at {checkpoints_dir}"}
+        
+        agent = ModelEnsembleAgent(
+            model_names=['swin', 'mobilenetv2', 'densenet169', 'efficientnetv2', 'maxvit'],
+            checkpoints_dir=checkpoints_dir,
+            num_classes=NUM_CLASSES,
+            class_names=CLASS_NAMES
+        )
+        
+        result = agent.run_ensemble(image_path)
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def run_educational_agent(diagnosis_result: Dict[str, Any], explanation_text: str = "") -> Dict[str, str]:
+    """Run the educational agent to translate diagnosis."""
+    try:
+        agent = EducationalAgent(doctor_name="your treating doctor")
+        
+        # Map ensemble result format to educational agent format
+        # Ensemble uses: ensemble_prediction, ensemble_confidence
+        # EducationalAgent expects: predicted_class, confidence_score
+        mapped_result = {
+            "predicted_class": diagnosis_result.get("ensemble_prediction", "Unknown"),
+            "confidence_score": diagnosis_result.get("ensemble_confidence", 0.0),
+            "fracture_detected": diagnosis_result.get("fracture_detected", True)
+        }
+        
+        result = agent.translate_to_layman_terms(mapped_result, explanation_text)
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def run_explainability_agent(diagnosis_result: Dict[str, Any]) -> str:
+    """Run the explainability agent to generate explanations."""
+    try:
+        agent = ExplainabilityAgent(class_names=CLASS_NAMES, body_part="bone")
+        
+        # Map ensemble result format to explainability agent format
+        # Ensemble uses: ensemble_prediction, ensemble_confidence
+        # ExplainabilityAgent expects: predicted_class, confidence_score
+        mapped_result = {
+            "predicted_class": diagnosis_result.get("ensemble_prediction", "Unknown"),
+            "confidence_score": diagnosis_result.get("ensemble_confidence", 0.0),
+            "fracture_detected": diagnosis_result.get("fracture_detected", True)
+        }
+        
+        # Generate a random heatmap for demonstration
+        heatmap = generate_random_heatmap()
+        
+        # Call with correct parameters
+        explanation = agent.generate_explanation(
+            diagnosis_result=mapped_result,
+            cam_array=heatmap
+        )
+        return explanation
+    except Exception as e:
+        return f"Error generating explanation: {str(e)}"
+
+
+def run_knowledge_agent(diagnosis: str, confidence: float) -> Dict[str, Any]:
+    """Run the knowledge agent to retrieve medical information."""
+    try:
+        agent = KnowledgeAgent(knowledge_base=MEDICAL_KNOWLEDGE_BASE)
+        result = agent.get_medical_summary(diagnosis, confidence)
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def run_complete_workflow(image_path: str) -> Dict[str, Any]:
+    """Run the complete workflow: Ensemble -> Education -> Knowledge."""
+    workflow_result = {
+        "ensemble_result": None,
+        "educational_result": None,
+        "knowledge_result": None,
+        "explanation_result": None
+    }
+    
+    try:
+        # 1. Run Ensemble Agent
+        ensemble_result = run_ensemble_agent(image_path)
+        if "error" in ensemble_result:
+            return {"error": f"Ensemble failed: {ensemble_result['error']}"}
+        
+        workflow_result["ensemble_result"] = ensemble_result
+        
+        # 2. Run Educational Agent
+        educational_result = run_educational_agent(ensemble_result)
+        workflow_result["educational_result"] = educational_result
+        
+        # 3. Run Explainability Agent
+        explanation = run_explainability_agent(ensemble_result)
+        workflow_result["explanation_result"] = explanation
+        
+        # 4. Run Knowledge Agent
+        diagnosis = ensemble_result.get("ensemble_prediction", "Unknown")
+        confidence = ensemble_result.get("ensemble_confidence", 0.0)
+        knowledge_result = run_knowledge_agent(diagnosis, confidence)
+        workflow_result["knowledge_result"] = knowledge_result
+        
+        return workflow_result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+# ============================================================================
+# --- 4. Streamlit UI ---
+# ============================================================================
+
+def main():
+    """Main Streamlit application."""
+    st.title("🦴 AI Medical Assistant for Fracture Detection & Diagnosis")
+    st.info("⚠️ **Research/Educational Use Only** - This system is not approved for clinical use without professional oversight.")
+    st.markdown("---")
+    
+    # Initialize session state
+    if "patient_context" not in st.session_state:
+        st.session_state.patient_context = {
+            "age": 45,
+            "gender": "Female",
+            "history": "No major past issues, but has mild osteoporosis."
+        }
+    
+    # Initialize workflow results storage
+    if "workflow_result" not in st.session_state:
+        st.session_state.workflow_result = None
+    
+    # --- Create Tabs ---
+    tab1, tab2, tab3, tab4, tab5 = st.tabs(
+        ["🏥 Single Agents", "⚙️ Complete Workflow", "💬 Patient Chat", "📋 Workflow Details", "ℹ️ About"]
+    )
+    
+    # ========================================================================
+    # --- TAB 1: Individual Agents ---
+    # ========================================================================
+    with tab1:
+        st.header("Run Individual Agents")
+        st.markdown("Test each agent independently with sample diagnosis data.")
+        
+        agent_choice = st.selectbox(
+            "Select an Agent",
+            ["Diagnostic Agent", "Ensemble Agent", "Educational Agent", "Explainability Agent", "Knowledge Agent"]
+        )
+        
+        # Create columns for layout
+        col1, col2 = st.columns([2, 1])
+        
+        with col1:
+            if agent_choice == "Diagnostic Agent":
+                st.subheader("🔍 Diagnostic Agent")
+                st.write("Runs a single model on an X-ray image to detect fractures.")
+                
+                image_file = st.file_uploader("Upload X-ray image", type=["jpg", "png", "jpeg"])
+                if image_file and st.button("Run Diagnostic Agent"):
+                    st.info("Note: Running this requires a valid model checkpoint at ./outputs/best_swin.pth")
+                    with st.spinner("Running diagnostic agent..."):
+                        image_path = save_uploaded_file(image_file)
+                        result = run_diagnostic_agent(image_path)
+                        st.json(result)
+            
+            elif agent_choice == "Ensemble Agent":
+                st.subheader("🎯 Ensemble Agent (5 Models)")
+                st.write("Combines predictions from multiple models for robust diagnosis.")
+                
+                image_file = st.file_uploader("Upload X-ray image", type=["jpg", "png", "jpeg"])
+                if image_file and st.button("Run Ensemble Agent"):
+                    st.info("Note: Running this requires model checkpoints in ./outputs/")
+                    with st.spinner("Running ensemble agent..."):
+                        image_path = save_uploaded_file(image_file)
+                        result = run_ensemble_agent(image_path)
+                        st.json(result)
+            
+            elif agent_choice == "Educational Agent":
+                st.subheader("📚 Educational Agent")
+                st.write("Translates technical diagnosis into patient-friendly language.")
+                
+                # Sample diagnosis for demo
+                sample_diagnosis = {
+                    "fracture_detected": True,
+                    "predicted_class": "Transverse",
+                    "confidence_score": 0.85,
+                    "severity_type": "Transverse"
+                }
+                
+                sample_explanation = "The bone shows a clear transverse break pattern."
+                
+                if st.button("Run Educational Agent (Demo)"):
+                    with st.spinner("Translating diagnosis..."):
+                        result = run_educational_agent(sample_diagnosis, sample_explanation)
+                        if isinstance(result, dict):
+                            for key, value in result.items():
+                                st.write(f"**{key}:**\n{value}")
+                        else:
+                            st.error(result)
+            
+            elif agent_choice == "Explainability Agent":
+                st.subheader("🎨 Explainability Agent")
+                st.write("Generates human-readable explanations of model predictions.")
+                
+                sample_diagnosis = {
+                    "predicted_class": "Greenstick",
+                    "confidence_score": 0.92,
+                    "fracture_detected": True
+                }
+                
+                if st.button("Run Explainability Agent (Demo)"):
+                    with st.spinner("Generating explanation..."):
+                        explanation = run_explainability_agent(sample_diagnosis)
+                        st.write(explanation)
+            
+            elif agent_choice == "Knowledge Agent":
+                st.subheader("🧠 Knowledge Agent")
+                st.write("Retrieves medical knowledge and guidelines for a diagnosis.")
+                
+                diagnosis_input = st.selectbox("Select Diagnosis", CLASS_NAMES)
+                confidence_input = st.slider("Confidence Score", 0.0, 1.0, 0.85)
+                
+                if st.button("Run Knowledge Agent"):
+                    with st.spinner("Retrieving medical knowledge..."):
+                        result = run_knowledge_agent(diagnosis_input, confidence_input)
+                        if isinstance(result, dict):
+                            st.json(result)
+                        else:
+                            st.error(result)
+    
+    # ========================================================================
+    # --- TAB 2: Complete Workflow ---
+    # ========================================================================
+    with tab2:
+        st.header("Complete Diagnosis Workflow")
+        st.markdown("Upload an X-ray image and run the complete diagnostic pipeline.")
+        
+        col1, col2 = st.columns([2, 1])
+        
+        with col1:
+            st.subheader("📤 Upload X-ray Image")
+            image_file = st.file_uploader("Upload X-ray image for full diagnosis", type=["jpg", "png", "jpeg"])
+            
+            if image_file:
+                st.image(image_file, caption="Uploaded Image", width='stretch')
+        
+        with col2:
+            st.subheader("👤 Patient Information")
+            age = st.number_input("Age", min_value=1, max_value=120, value=st.session_state.patient_context["age"])
+            gender = st.selectbox("Gender", ["Male", "Female", "Other"], 
+                                 index=0 if st.session_state.patient_context["gender"] == "Male" else 
+                                       1 if st.session_state.patient_context["gender"] == "Female" else 2)
+            history = st.text_area("Medical History", value=st.session_state.patient_context["history"])
+            
+            st.session_state.patient_context = {"age": age, "gender": gender, "history": history}
+        
+        if image_file and st.button("🚀 Run Complete Workflow", key="workflow"):
+            st.info("Note: Running this requires all model checkpoints in ./outputs/")
+            with st.spinner("Running complete diagnostic workflow..."):
+                image_path = save_uploaded_file(image_file)
+                workflow_result = run_complete_workflow(image_path)
+                
+                # Store workflow result in session state for use in other tabs
+                st.session_state.workflow_result = workflow_result
+                
+                if "error" in workflow_result:
+                    st.error(f"❌ Error: {workflow_result['error']}")
+                else:
+                    # Display results
+                    st.success("✅ Workflow completed successfully!")
+                    
+                    # Ensemble Results
+                    if workflow_result["ensemble_result"]:
+                        st.subheader("1️⃣ Ensemble Agent Results")
+                        ensemble = workflow_result["ensemble_result"]
+                        col1, col2, col3 = st.columns(3)
+                        col1.metric("Prediction", ensemble.get("ensemble_prediction", "N/A"))
+                        col2.metric("Confidence", f"{ensemble.get('ensemble_confidence', 0):.2%}")
+                        col3.metric("Fracture Detected", "Yes" if ensemble.get("fracture_detected") else "No")
+                    
+                    # Educational Results
+                    if workflow_result["educational_result"]:
+                        st.subheader("2️⃣ Patient-Friendly Summary")
+                        educational = workflow_result["educational_result"]
+                        for key, value in educational.items():
+                            st.write(f"**{key}:**\n{value}")
+                    
+                    # Explainability Results
+                    if workflow_result["explanation_result"]:
+                        st.subheader("3️⃣ Technical Explanation")
+                        st.write(workflow_result["explanation_result"])
+                    
+                    # Knowledge Results
+                    if workflow_result["knowledge_result"]:
+                        st.subheader("4️⃣ Medical Knowledge Base")
+                        st.json(workflow_result["knowledge_result"])
+    
+    # ========================================================================
+    # --- TAB 3: Patient Chat (Hugging Face) ---
+    # ========================================================================
+    with tab3:
+        st.header("💬 Patient Q&A with AI Assistant")
+        st.markdown("Ask questions about your fracture diagnosis using Hugging Face Inference API")
+        
+        # Check if workflow has been run
+        if st.session_state.workflow_result is None or "error" in st.session_state.workflow_result:
+            st.info("ℹ️ Please run the 'Complete Workflow' first to generate a diagnosis for the chat feature.")
+        else:
+            # Check HF API configuration
+            if not HF_API_KEY:
+                st.error(
+                    "❌ Hugging Face API key not configured. "
+                    "Please add your HUGGINGFACE_API_KEY to Streamlit Secrets."
+                )
+                st.markdown("""
+                ### How to set up Hugging Face API:
+                1. Get your API key from https://huggingface.co/settings/tokens
+                2. In Streamlit Cloud, go to Settings > Secrets
+                3. Add: `HUGGINGFACE_API_KEY = "hf_your_token_here"`
+                4. Refresh the app
+                """)
+            else:
+                # Build medical summary from workflow results
+                ensemble_result = st.session_state.workflow_result.get("ensemble_result", {})
+                knowledge_result = st.session_state.workflow_result.get("knowledge_result", {})
+                
+                diagnosis = ensemble_result.get("ensemble_prediction", "Unknown")
+                confidence = ensemble_result.get("ensemble_confidence", 0.0)
+                
+                # Create medical summary from knowledge base
+                medical_summary = {
+                    "Diagnosis": diagnosis,
+                    "Ensemble_Confidence": f"{confidence:.2f}",
+                    "Type": knowledge_result.get("Type", "Unknown fracture type"),
+                    "Severity": knowledge_result.get("Severity", "Unknown"),
+                    "Guidelines": knowledge_result.get("Guidelines", [])
+                }
+                
+                try:
+                    agent = PatientInteractionAgent(medical_summary, st.session_state.patient_context)
+                    
+                    # Initialize chat history with diagnosis info
+                    if "messages" not in st.session_state:
+                        st.session_state.messages = []
+                        st.session_state.messages.append({
+                            "role": "assistant",
+                            "content": f"Hello! I'm your AI medical assistant. I've reviewed your diagnosis: **{medical_summary['Diagnosis']}** (Confidence: {medical_summary['Ensemble_Confidence']}). How can I help answer your questions?"
+                        })
+                    
+                    # Display chat messages
+                    for message in st.session_state.messages:
+                        with st.chat_message(message["role"]):
+                            st.markdown(message["content"])
+                    
+                    # Accept user input
+                    if prompt := st.chat_input("Ask a question about your diagnosis..."):
+                        st.session_state.messages.append({"role": "user", "content": prompt})
+                        with st.chat_message("user"):
+                            st.markdown(prompt)
+                        
+                        with st.chat_message("assistant"):
+                            with st.spinner("🤖 Consulting Mistral 7B via Hugging Face..."):
+                                response = agent.get_response(prompt)
+                                st.markdown(response)
+                        
+                        st.session_state.messages.append({"role": "assistant", "content": response})
+                
+                except ValueError as e:
+                    st.error(str(e))
+                except Exception as e:
+                    st.error(f"❌ Error initializing chat agent: {str(e)}")
+    
+    # ========================================================================
+    # --- TAB 4: Workflow Details ---
+    # ========================================================================
+    with tab4:
+        st.header("📋 Workflow Execution Details")
+        
+        if st.session_state.workflow_result is None:
+            st.info("ℹ️ No workflow results available. Please run a workflow first.")
+        else:
+            if "error" in st.session_state.workflow_result:
+                st.error(f"Workflow Error: {st.session_state.workflow_result['error']}")
+            else:
+                st.success("Workflow executed successfully!")
+                st.json(st.session_state.workflow_result)
+    
+    # ========================================================================
+    # --- TAB 5: About ---
+    # ========================================================================
+    with tab5:
+        st.header("ℹ️ About This Application")
+        st.markdown("""
+        ### 🦴 AI-Powered Fracture Detection System
+        
+        This application uses advanced deep learning models to detect and classify fractures from X-ray images.
+        
+        **Features:**
+        - **Multi-Model Ensemble:** Combines 5 different architectures (Swin, MobileNetV2, DenseNet, EfficientNet, MaxViT)
+        - **Explainability:** Generates human-readable explanations for predictions
+        - **Patient Education:** Translates medical terminology into patient-friendly language
+        - **AI Chatbot:** Ask questions about your diagnosis powered by Mistral 7B via Hugging Face
+        
+        **Models Used:**
+        - Swin Transformer
+        - MobileNetV2
+        - DenseNet169
+        - EfficientNetV2
+        - MaxViT
+        
+        **Fracture Types Detected:**
+        """)
+        
+        for i, fracture_type in enumerate(CLASS_NAMES, 1):
+            st.write(f"{i}. {fracture_type}")
+        
+        st.markdown("""
+        ### ⚠️ Important Disclaimer
+        This system is for **research and educational purposes only**. 
+        It is **NOT approved for clinical use** without professional medical oversight.
+        Always consult with a qualified healthcare professional for medical diagnosis.
+        
+        ### 🔧 Technology Stack
+        - **Frontend:** Streamlit
+        - **ML Models:** PyTorch
+        - **AI Assistant:** Hugging Face Inference API (Mistral 7B)
+        - **Deployment:** Streamlit Cloud
+        
+        ### 📞 Contact & Support
+        For issues or questions, please contact the development team.
+        """)
+
+
+if __name__ == "__main__":
+    main()

apps/patient_chat_app_local.py

@@ -0,0 +1,663 @@
+"""
+Streamlit-based Patient Chat Application for Fracture Detection and Diagnosis.
+
+Supports:
+1. Running individual agents (Diagnostic, Educational, Explainability, Knowledge)
+2. Running the complete workflow
+3. LLM-based Q&A for patient education
+"""
+
+import os
+import sys
+import streamlit as st
+import requests
+import json
+import numpy as np
+from typing import Dict, Any, List
+from pathlib import Path
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+# --- Import the Agents ---
+from src.agents.diagnostic_agent import DiagnosticAgent
+from src.agents.educational_agent import EducationalAgent
+from src.agents.explain_agent import ExplainabilityAgent, generate_random_heatmap, calculate_heatmap_centroid
+from src.agents.knowledge_agent import KnowledgeAgent, MEDICAL_KNOWLEDGE_BASE
+from src.agents.cross_validation_agent import ModelEnsembleAgent
+from src.utils import get_device
+
+# --- Configuration for Ollama ---
+# Support both localhost and host.docker.internal for Docker deployments
+OLLAMA_ENDPOINT = os.getenv("OLLAMA_ENDPOINT", os.getenv("OLLAMA_HOST", "http://localhost:11434") + "/api/generate")
+OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")  # Ensure you have pulled this model using 'ollama pull llama3'
+OLLAMA_CHECK_URL = os.getenv("OLLAMA_HOST", "http://localhost:11434")
+
+# --- Constants ---
+CLASS_NAMES = ["Comminuted", "Greenstick", "Healthy", "Oblique", 
+               "Oblique Displaced", "Spiral", "Transverse", "Transverse Displaced"]
+NUM_CLASSES = len(CLASS_NAMES)
+IMG_SIZE = 224
+
+# --- Page Configuration ---
+st.set_page_config(
+    page_title="🦴 Fracture Detection AI System",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# --- Custom CSS for Better UI ---
+st.markdown("""
+<style>
+    .stTabs [data-baseweb="tab-list"] button {
+        font-size: 16px;
+        font-weight: bold;
+    }
+    .section-header {
+        font-size: 20px;
+        font-weight: bold;
+        margin-top: 20px;
+        margin-bottom: 10px;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+
+# ============================================================================
+# --- 1. Ollama-based Patient Interaction Agent ---
+# ============================================================================
+
+class PatientInteractionAgent:
+    """
+    Handles communication with the local Llama 3 model via the Ollama API endpoint.
+    """
+    def __init__(self, medical_summary: Dict[str, Any], patient_history: Dict[str, Any]):
+        """Initialize the agent with medical context."""
+        # --- Connection Check ---
+        try:
+            response = requests.get(OLLAMA_CHECK_URL, timeout=5)
+            if response.status_code != 200:
+                raise ConnectionError("Ollama server is not running or accessible.")
+        except requests.exceptions.ConnectionError:
+            raise ConnectionError("Ollama server is not running. Please start Ollama.")
+
+        self.medical_summary = medical_summary
+        self.patient_history = patient_history
+        self.system_prompt = self._build_system_prompt()
+
+    def _build_system_prompt(self) -> str:
+        """Creates a detailed instruction set for the LLM (RAG Context)."""
+        guidelines = "\n- ".join(self.medical_summary.get('Guidelines', ["No specific guidelines available."]))
+
+        return f"""
+You are a highly compassionate, clear, and professional medical assistant. Your goal is to answer patient questions
+in natural language based ONLY on the following diagnostic information and patient history.
+
+RULES:
+1. Maintain a reassuring, non-technical, and empathetic tone suitable for a patient.
+2. Keep answers concise and address the patient's underlying concern.
+3. ALWAYS conclude your answer by advising the patient to consult their orthopedic specialist or doctor.
+
+--- DIAGNOSTIC INFORMATION ---
+Diagnosis: {self.medical_summary.get('Diagnosis')} (Confidence: {self.medical_summary.get('Ensemble_Confidence')})
+Definition: {self.medical_summary.get('Type')}
+Severity: {self.medical_summary.get('Severity')}
+Treatment Guidelines:
+{guidelines}
+
+--- PATIENT HISTORY ---
+Age: {self.patient_history.get('age')}
+Gender: {self.patient_history.get('gender')}
+Medical History: {self.patient_history.get('history')}
+"""
+
+    def get_response(self, query: str) -> str:
+        """Sends the user query to the Llama 3 model via Ollama."""
+        full_prompt = f"{self.system_prompt}\n\nPATIENT QUERY: {query}"
+
+        payload = {
+            "model": OLLAMA_MODEL,
+            "prompt": full_prompt,
+            "stream": False,
+            "options": {"temperature": 0.1}
+        }
+
+        try:
+            response = requests.post(OLLAMA_ENDPOINT, json=payload, timeout=300)
+            response.raise_for_status()
+            data = response.json()
+            return data.get("response", "Error: Could not extract response from Ollama.")
+        except requests.exceptions.RequestException as e:
+            return f"Error communicating with Ollama: {e}"
+        except Exception as e:
+            return f"An unexpected error occurred: {e}"
+
+
+# ============================================================================
+# --- 2. Helper Functions ---
+# ============================================================================
+
+def save_uploaded_file(uploaded_file) -> str:
+    """Save uploaded file to temp location and return path."""
+    if uploaded_file is None:
+        return None
+    
+    try:
+        import tempfile
+        # Create a temporary file
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file:
+            tmp_file.write(uploaded_file.getbuffer())
+            return tmp_file.name
+    except Exception as e:
+        st.error(f"Error saving file: {e}")
+        return None
+
+
+# ============================================================================
+# --- 3. Workflow Functions ---
+# ============================================================================
+
+def run_diagnostic_agent(image_path: str) -> Dict[str, Any]:
+    """Run the diagnostic agent on an image."""
+    try:
+        # Placeholder checkpoint path - in production, use actual model checkpoint
+        checkpoint_path = "./outputs/best_swin.pth"
+        
+        if not os.path.exists(checkpoint_path):
+            return {"error": f"Checkpoint not found at {checkpoint_path}"}
+        
+        agent = DiagnosticAgent(
+            checkpoint_path=checkpoint_path,
+            model_name='swin',
+            num_classes=NUM_CLASSES,
+            img_size=IMG_SIZE,
+            class_names=CLASS_NAMES
+        )
+        
+        result = agent.run_diagnosis(image_path)
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def run_ensemble_agent(image_path: str) -> Dict[str, Any]:
+    """Run the ensemble agent on an image."""
+    try:
+        checkpoints_dir = "./outputs"
+        
+        if not os.path.exists(checkpoints_dir):
+            return {"error": f"Checkpoints directory not found at {checkpoints_dir}"}
+        
+        agent = ModelEnsembleAgent(
+            model_names=['swin', 'mobilenetv2', 'densenet169', 'efficientnetv2', 'maxvit'],
+            checkpoints_dir=checkpoints_dir,
+            num_classes=NUM_CLASSES,
+            class_names=CLASS_NAMES
+        )
+        
+        result = agent.run_ensemble(image_path)
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def run_educational_agent(diagnosis_result: Dict[str, Any], explanation_text: str = "") -> Dict[str, str]:
+    """Run the educational agent to translate diagnosis."""
+    try:
+        agent = EducationalAgent(doctor_name="your treating doctor")
+        
+        # Map ensemble result format to educational agent format
+        # Ensemble uses: ensemble_prediction, ensemble_confidence
+        # EducationalAgent expects: predicted_class, confidence_score
+        mapped_result = {
+            "predicted_class": diagnosis_result.get("ensemble_prediction", "Unknown"),
+            "confidence_score": diagnosis_result.get("ensemble_confidence", 0.0),
+            "fracture_detected": diagnosis_result.get("fracture_detected", True)
+        }
+        
+        result = agent.translate_to_layman_terms(mapped_result, explanation_text)
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def run_explainability_agent(diagnosis_result: Dict[str, Any]) -> str:
+    """Run the explainability agent to generate explanations."""
+    try:
+        agent = ExplainabilityAgent(class_names=CLASS_NAMES, body_part="bone")
+        
+        # Map ensemble result format to explainability agent format
+        # Ensemble uses: ensemble_prediction, ensemble_confidence
+        # ExplainabilityAgent expects: predicted_class, confidence_score
+        mapped_result = {
+            "predicted_class": diagnosis_result.get("ensemble_prediction", "Unknown"),
+            "confidence_score": diagnosis_result.get("ensemble_confidence", 0.0),
+            "fracture_detected": diagnosis_result.get("fracture_detected", True)
+        }
+        
+        # Generate a random heatmap for demonstration
+        heatmap = generate_random_heatmap()
+        
+        # Call with correct parameters
+        explanation = agent.generate_explanation(
+            diagnosis_result=mapped_result,
+            cam_array=heatmap
+        )
+        return explanation
+    except Exception as e:
+        return f"Error generating explanation: {str(e)}"
+
+
+def run_knowledge_agent(diagnosis: str, confidence: float) -> Dict[str, Any]:
+    """Run the knowledge agent to retrieve medical information."""
+    try:
+        agent = KnowledgeAgent(knowledge_base=MEDICAL_KNOWLEDGE_BASE)
+        result = agent.get_medical_summary(diagnosis, confidence)
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def run_complete_workflow(image_path: str) -> Dict[str, Any]:
+    """Run the complete workflow: Ensemble -> Education -> Knowledge."""
+    workflow_result = {
+        "ensemble_result": None,
+        "educational_result": None,
+        "knowledge_result": None,
+        "explanation_result": None
+    }
+    
+    try:
+        # 1. Run Ensemble Agent
+        ensemble_result = run_ensemble_agent(image_path)
+        if "error" in ensemble_result:
+            return {"error": f"Ensemble failed: {ensemble_result['error']}"}
+        
+        workflow_result["ensemble_result"] = ensemble_result
+        
+        # 2. Run Educational Agent
+        educational_result = run_educational_agent(ensemble_result)
+        workflow_result["educational_result"] = educational_result
+        
+        # 3. Run Explainability Agent
+        explanation = run_explainability_agent(ensemble_result)
+        workflow_result["explanation_result"] = explanation
+        
+        # 4. Run Knowledge Agent
+        diagnosis = ensemble_result.get("ensemble_prediction", "Unknown")
+        confidence = ensemble_result.get("ensemble_confidence", 0.0)
+        knowledge_result = run_knowledge_agent(diagnosis, confidence)
+        workflow_result["knowledge_result"] = knowledge_result
+        
+        return workflow_result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+# ============================================================================
+# --- 3. Streamlit UI ---
+# ============================================================================
+
+def main():
+    """Main Streamlit application."""
+    st.title("🦴 AI Medical Assistant for Fracture Detection & Diagnosis")
+    st.info("⚠️ **Research/Educational Use Only** - This system is not approved for clinical use without professional oversight.")
+    st.markdown("---")
+    
+    # Initialize session state
+    if "patient_context" not in st.session_state:
+        st.session_state.patient_context = {
+            "age": 45,
+            "gender": "Female",
+            "history": "No major past issues, but has mild osteoporosis."
+        }
+    
+    # Initialize workflow results storage
+    if "workflow_result" not in st.session_state:
+        st.session_state.workflow_result = None
+    
+    # --- Create Tabs ---
+    tab1, tab2, tab3, tab4, tab5 = st.tabs(
+        ["🏥 Single Agents", "⚙️ Complete Workflow", "💬 Patient Chat", "📋 Workflow Details", "ℹ️ About"]
+    )
+    
+    # ========================================================================
+    # --- TAB 1: Individual Agents ---
+    # ========================================================================
+    with tab1:
+        st.header("Run Individual Agents")
+        st.markdown("Test each agent independently with sample diagnosis data.")
+        
+        agent_choice = st.selectbox(
+            "Select an Agent",
+            ["Diagnostic Agent", "Ensemble Agent", "Educational Agent", "Explainability Agent", "Knowledge Agent"]
+        )
+        
+        # Create columns for layout
+        col1, col2 = st.columns([2, 1])
+        
+        with col1:
+            if agent_choice == "Diagnostic Agent":
+                st.subheader("🔍 Diagnostic Agent")
+                st.write("Runs a single model on an X-ray image to detect fractures.")
+                
+                image_file = st.file_uploader("Upload X-ray image", type=["jpg", "png", "jpeg"])
+                if image_file and st.button("Run Diagnostic Agent"):
+                    st.info("Note: Running this requires a valid model checkpoint at ./outputs/best_swin.pth")
+                    with st.spinner("Running diagnostic agent..."):
+                        image_path = save_uploaded_file(image_file)
+                        result = run_diagnostic_agent(image_path)
+                        st.json(result)
+            
+            elif agent_choice == "Ensemble Agent":
+                st.subheader("🎯 Ensemble Agent (5 Models)")
+                st.write("Combines predictions from multiple models for robust diagnosis.")
+                
+                image_file = st.file_uploader("Upload X-ray image", type=["jpg", "png", "jpeg"])
+                if image_file and st.button("Run Ensemble Agent"):
+                    st.info("Note: Running this requires model checkpoints in ./outputs/")
+                    with st.spinner("Running ensemble agent..."):
+                        image_path = save_uploaded_file(image_file)
+                        result = run_ensemble_agent(image_path)
+                        st.json(result)
+            
+            elif agent_choice == "Educational Agent":
+                st.subheader("📚 Educational Agent")
+                st.write("Translates technical diagnosis into patient-friendly language.")
+                
+                # Sample diagnosis for demo
+                sample_diagnosis = {
+                    "fracture_detected": True,
+                    "predicted_class": "Transverse",
+                    "confidence_score": 0.85,
+                    "severity_type": "Transverse"
+                }
+                
+                sample_explanation = "The bone shows a clear transverse break pattern."
+                
+                if st.button("Run Educational Agent (Demo)"):
+                    with st.spinner("Translating diagnosis..."):
+                        result = run_educational_agent(sample_diagnosis, sample_explanation)
+                        if isinstance(result, dict):
+                            for key, value in result.items():
+                                st.write(f"**{key}:**\n{value}")
+                        else:
+                            st.error(result)
+            
+            elif agent_choice == "Explainability Agent":
+                st.subheader("🎨 Explainability Agent")
+                st.write("Generates human-readable explanations of model predictions.")
+                
+                sample_diagnosis = {
+                    "predicted_class": "Greenstick",
+                    "confidence_score": 0.92,
+                    "fracture_detected": True
+                }
+                
+                if st.button("Run Explainability Agent (Demo)"):
+                    with st.spinner("Generating explanation..."):
+                        explanation = run_explainability_agent(sample_diagnosis)
+                        st.write(explanation)
+            
+            elif agent_choice == "Knowledge Agent":
+                st.subheader("🧠 Knowledge Agent")
+                st.write("Retrieves medical knowledge and guidelines for a diagnosis.")
+                
+                diagnosis_input = st.selectbox("Select Diagnosis", CLASS_NAMES)
+                confidence_input = st.slider("Confidence Score", 0.0, 1.0, 0.85)
+                
+                if st.button("Run Knowledge Agent"):
+                    with st.spinner("Retrieving medical knowledge..."):
+                        result = run_knowledge_agent(diagnosis_input, confidence_input)
+                        if isinstance(result, dict):
+                            st.json(result)
+                        else:
+                            st.error(result)
+    
+    # ========================================================================
+    # --- TAB 2: Complete Workflow ---
+    # ========================================================================
+    with tab2:
+        st.header("Complete Diagnosis Workflow")
+        st.markdown("Upload an X-ray image and run the complete diagnostic pipeline.")
+        
+        col1, col2 = st.columns([2, 1])
+        
+        with col1:
+            st.subheader("📤 Upload X-ray Image")
+            image_file = st.file_uploader("Upload X-ray image for full diagnosis", type=["jpg", "png", "jpeg"])
+            
+            if image_file:
+                st.image(image_file, caption="Uploaded Image", width='stretch')
+        
+        with col2:
+            st.subheader("👤 Patient Information")
+            age = st.number_input("Age", min_value=1, max_value=120, value=st.session_state.patient_context["age"])
+            gender = st.selectbox("Gender", ["Male", "Female", "Other"], 
+                                 index=0 if st.session_state.patient_context["gender"] == "Male" else 
+                                       1 if st.session_state.patient_context["gender"] == "Female" else 2)
+            history = st.text_area("Medical History", value=st.session_state.patient_context["history"])
+            
+            st.session_state.patient_context = {"age": age, "gender": gender, "history": history}
+        
+        if image_file and st.button("🚀 Run Complete Workflow", key="workflow"):
+            st.info("Note: Running this requires all model checkpoints in ./outputs/")
+            with st.spinner("Running complete diagnostic workflow..."):
+                image_path = save_uploaded_file(image_file)
+                workflow_result = run_complete_workflow(image_path)
+                
+                # Store workflow result in session state for use in other tabs
+                st.session_state.workflow_result = workflow_result
+                
+                if "error" in workflow_result:
+                    st.error(f"❌ Error: {workflow_result['error']}")
+                else:
+                    # Display results
+                    st.success("✅ Workflow completed successfully!")
+                    
+                    # Ensemble Results
+                    if workflow_result["ensemble_result"]:
+                        st.subheader("1️⃣ Ensemble Agent Results")
+                        ensemble = workflow_result["ensemble_result"]
+                        col1, col2, col3 = st.columns(3)
+                        col1.metric("Prediction", ensemble.get("ensemble_prediction", "N/A"))
+                        col2.metric("Confidence", f"{ensemble.get('ensemble_confidence', 0):.2%}")
+                        col3.metric("Fracture Detected", "Yes" if ensemble.get("fracture_detected") else "No")
+                    
+                    # Educational Results
+                    if workflow_result["educational_result"]:
+                        st.subheader("2️⃣ Patient-Friendly Summary")
+                        educational = workflow_result["educational_result"]
+                        for key, value in educational.items():
+                            st.write(f"**{key}:**\n{value}")
+                    
+                    # Explainability Results
+                    if workflow_result["explanation_result"]:
+                        st.subheader("3️⃣ Technical Explanation")
+                        st.write(workflow_result["explanation_result"])
+                    
+                    # Knowledge Results
+                    if workflow_result["knowledge_result"]:
+                        st.subheader("4️⃣ Medical Knowledge Base")
+                        st.json(workflow_result["knowledge_result"])
+    
+    # ========================================================================
+    # --- TAB 3: Patient Chat ---
+    # ========================================================================
+    with tab3:
+        st.header("💬 Patient Q&A with AI Assistant")
+        st.markdown("Ask questions about your fracture diagnosis (requires Ollama running)")
+        
+        # Check if workflow has been run
+        if st.session_state.workflow_result is None or "error" in st.session_state.workflow_result:
+            st.info("ℹ️ Please run the 'Complete Workflow' first to generate a diagnosis for the chat feature.")
+        else:
+            # Check for Ollama availability
+            ollama_available = False
+            try:
+                response = requests.get(OLLAMA_CHECK_URL, timeout=2)
+                ollama_available = response.status_code == 200
+            except:
+                ollama_available = False
+            
+            if not ollama_available:
+                st.warning("⚠️ Ollama server is not running. Please start Ollama to use the chat feature.")
+                st.info("Download Ollama from https://ollama.ai and run: ollama pull llama3")
+            else:
+                # Build medical summary from workflow results
+                ensemble_result = st.session_state.workflow_result.get("ensemble_result", {})
+                knowledge_result = st.session_state.workflow_result.get("knowledge_result", {})
+                
+                diagnosis = ensemble_result.get("ensemble_prediction", "Unknown")
+                confidence = ensemble_result.get("ensemble_confidence", 0.0)
+                
+                # Create medical summary from knowledge base
+                medical_summary = {
+                    "Diagnosis": diagnosis,
+                    "Ensemble_Confidence": f"{confidence:.2f}",
+                    "Type": knowledge_result.get("Type", "Unknown fracture type"),
+                    "Severity": knowledge_result.get("Severity", "Unknown"),
+                    "Guidelines": knowledge_result.get("Guidelines", [])
+                }
+                
+                try:
+                    agent = PatientInteractionAgent(medical_summary, st.session_state.patient_context)
+                    
+                    # Initialize chat history with diagnosis info
+                    if "messages" not in st.session_state:
+                        st.session_state.messages = []
+                        st.session_state.messages.append({
+                            "role": "assistant",
+                            "content": f"Hello! I'm your AI medical assistant. I've reviewed your diagnosis: **{medical_summary['Diagnosis']}** (Confidence: {medical_summary['Ensemble_Confidence']}). How can I help answer your questions?"
+                        })
+                    
+                    # Display chat messages
+                    for message in st.session_state.messages:
+                        with st.chat_message(message["role"]):
+                            st.markdown(message["content"])
+                    
+                    # Accept user input
+                    if prompt := st.chat_input("Ask a question about your diagnosis..."):
+                        st.session_state.messages.append({"role": "user", "content": prompt})
+                        with st.chat_message("user"):
+                            st.markdown(prompt)
+                        
+                        with st.chat_message("assistant"):
+                            with st.spinner(f"Asking {OLLAMA_MODEL}..."):
+                                response = agent.get_response(prompt)
+                                st.markdown(response)
+                        
+                        st.session_state.messages.append({"role": "assistant", "content": response})
+                
+                except ConnectionError as e:
+                    st.error(f"❌ Connection Error: {e}")
+                except Exception as e:
+                    st.error(f"❌ Error: {e}")
+    
+    # ========================================================================
+    # --- TAB 4: Workflow Details ---
+    # ========================================================================
+    with tab4:
+        st.header("📋 System Architecture & Workflow")
+        
+        st.subheader("1. Ensemble Agent (Cross-Validation)")
+        st.write("""
+        - **Purpose**: Combines predictions from 5 different deep learning models
+        - **Models**: Swin, MobileNetV2, DenseNet169, EfficientNetV2, MaxViT
+        - **Output**: Ensemble prediction with confidence score
+        - **Benefit**: More robust and reliable predictions than single model
+        """)
+        
+        st.subheader("2. Educational Agent")
+        st.write("""
+        - **Purpose**: Translates technical diagnosis into patient-friendly language
+        - **Input**: Diagnosis result from ensemble
+        - **Output**: 
+          - Patient summary
+          - Severity assessment in simple terms
+          - Next steps and action plan
+        """)
+        
+        st.subheader("3. Explainability Agent")
+        st.write("""
+        - **Purpose**: Generates visual and textual explanations of predictions
+        - **Input**: Diagnosis result and Grad-CAM heatmap
+        - **Output**: Human-readable explanation of what the model "saw"
+        """)
+        
+        st.subheader("4. Knowledge Agent")
+        st.write("""
+        - **Purpose**: Retrieves medical knowledge for each diagnosis
+        - **Input**: Final diagnosis and confidence
+        - **Output**:
+          - Medical definition
+          - ICD code
+          - Treatment guidelines
+          - Severity level
+        """)
+        
+        st.markdown("---")
+        st.markdown("### Workflow Pipeline")
+        st.markdown("""
+        ```
+        X-ray Image
+            ↓
+        [Ensemble Agent] → Diagnosis + Confidence
+            ↓
+        [Educational Agent] → Patient-Friendly Summary
+            ↓
+        [Explainability Agent] → Visual Explanation
+            ↓
+        [Knowledge Agent] → Medical Guidelines
+            ↓
+        Patient Report
+        ```
+        """)
+    
+    # ========================================================================
+    # --- TAB 5: About ---
+    # ========================================================================
+    with tab5:
+        st.header("ℹ️ About This System")
+        
+        st.markdown("""
+        ## MedAI - Explainable Fracture Detection
+        
+        This application demonstrates an AI-powered medical diagnosis system designed to assist
+        healthcare professionals in fracture detection and patient education.
+        
+        ### Features:
+        - 🎯 **Ensemble Learning**: 5 deep learning models for robust predictions
+        - 📚 **Patient Education**: Automatic translation of technical diagnoses
+        - 🎨 **Explainability**: Visual and textual explanations of AI decisions
+        - 🧠 **Knowledge Integration**: Evidence-based medical guidelines
+        - 💬 **LLM Integration**: Natural language Q&A with Llama 3
+        
+        ### Supported Fracture Types:
+        - Comminuted
+        - Greenstick
+        - Oblique
+        - Oblique Displaced
+        - Spiral
+        - Transverse
+        - Transverse Displaced
+        - Healthy (No fracture)
+        
+        ### Technical Stack:
+        - **Deep Learning**: PyTorch with timm models
+        - **Frontend**: Streamlit
+        - **LLM**: Llama 3 via Ollama
+        - **Explainability**: Grad-CAM + Centroid analysis
+        
+        ### Disclaimer:
+        This system is for educational and research purposes. It should not be used
+        for actual medical diagnosis without proper clinical validation and oversight.
+        Always consult with qualified medical professionals for diagnosis and treatment.
+        """)
+        
+        st.markdown("---")
+        st.info("📧 For more information, visit the project repository on GitHub.")
+
+
+if __name__ == "__main__":
+    main()

outputs/best_densenet169.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3686669a0348a9043c9a697ec4a13463b4acc07649cb682cf44814ea6c3ac085
+size 195498793

outputs/best_efficientnetv2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2979a53289c91ab60e2697141952c0ecb66174c8d3ae998ccb16b8065b5fc93f
+size 195498793

outputs/best_maxvit.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4819e32d99287b04501f7e52fd80e9eff86a92584eeef7c394f70729180538a9
+size 195498793

outputs/best_mobilenetv2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efea091ddf0c43ce7cbca8a7eefe450974810d537984ae9521dc44ed3aa9055a
+size 195498793

outputs/best_swin.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:095f19e8c238fda959f7412f6c39a4f43362ea03178efb1535f93dc8386763f7
+size 195498793

requirements-prod.txt

@@ -0,0 +1,21 @@
+# Core ML/DL
+torch==2.9.0
+torchvision==0.24.0
+timm==1.0.20
+
+# Web Framework
+streamlit==1.51.0
+
+# Data & ML Tools
+numpy
+pandas
+scikit-learn==1.7.2
+Pillow
+opencv-python
+
+# API & HTTP
+requests
+
+# Utilities
+tqdm
+pyyaml

src/__init__.py

@@ -0,0 +1,24 @@
+"""
+Source code for MedAI Fracture Detection System.
+"""
+
+__version__ = "1.0.0"
+
+# Import key utilities for easy access
+from .utils import (
+    get_device,
+    require_mps,
+    DEVICE,
+    get_model,
+    get_transforms,
+    FractureDataset
+)
+
+__all__ = [
+    'get_device',
+    'require_mps',
+    'DEVICE',
+    'get_model',
+    'get_transforms',
+    'FractureDataset'
+]

src/agents/__init__.py

@@ -0,0 +1,19 @@
+"""
+Agent modules for fracture detection and diagnosis system.
+"""
+
+from .diagnostic_agent import DiagnosticAgent
+from .explain_agent import generate_random_heatmap, calculate_heatmap_centroid
+from .educational_agent import EducationalAgent
+from .knowledge_agent import KnowledgeAgent, MEDICAL_KNOWLEDGE_BASE
+from .cross_validation_agent import ModelEnsembleAgent
+
+__all__ = [
+    'DiagnosticAgent',
+    'generate_random_heatmap',
+    'calculate_heatmap_centroid',
+    'EducationalAgent',
+    'KnowledgeAgent',
+    'MEDICAL_KNOWLEDGE_BASE',
+    'ModelEnsembleAgent'
+]

src/agents/cross_validation_agent.py

@@ -0,0 +1,183 @@
+import os
+import sys
+import argparse
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+import numpy as np
+from PIL import Image
+from typing import List, Dict, Any
+import timm 
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+from src.utils import get_device, get_model, get_transforms
+
+# ----------------------------------------------------------------------
+# --- Global Variables ---
+# ----------------------------------------------------------------------
+
+DEVICE = get_device()
+IMG_SIZE = 224
+
+# ----------------------------------------------------------------------
+# --- Model Ensemble Agent Core (with all fixes) ---
+# ----------------------------------------------------------------------
+
+class ModelEnsembleAgent:
+    def __init__(self, model_names: List[str], checkpoints_dir: str, num_classes: int, class_names: List[str]):
+        self.models = {}
+        self.model_names = model_names
+        self.num_classes = num_classes
+        self.class_names = class_names
+        self.transforms = get_transforms('val', IMG_SIZE)
+        
+        self.device = DEVICE
+        self._load_all_models(checkpoints_dir)
+
+    def _load_all_models(self, checkpoints_dir: str):
+        """Loads all specified model checkpoints with strict=False fallback."""
+        print(f"Loading {len(self.model_names)} models from {checkpoints_dir} on {self.device}...")
+        
+        for name in self.model_names:
+            
+            # FIX: Corrected file naming convention (best_modelname.pth)
+            checkpoint_path = os.path.join(checkpoints_dir, f"best_{name}.pth")
+            
+            print(f"  Attempting to load {name} from expected path: {checkpoint_path}...")
+            
+            try:
+                model = get_model(name, self.num_classes, pretrained=False).to(self.device)
+                
+                checkpoint = torch.load(checkpoint_path, map_location=self.device)
+                state_dict = checkpoint.get('model_state_dict', checkpoint)
+                
+                # FIX: Filter out incompatible head layers that have size mismatches
+                # This handles cases where checkpoint was trained with different head architecture
+                model_state = model.state_dict()
+                filtered_state_dict = {}
+                for key, value in state_dict.items():
+                    if key in model_state and model_state[key].shape == value.shape:
+                        filtered_state_dict[key] = value
+                    elif key not in model_state:
+                        # Key doesn't exist in current model, skip it
+                        pass
+                    else:
+                        # Shape mismatch, skip this layer (usually head layers)
+                        print(f"    (Skipping layer '{key}' due to shape mismatch: {value.shape} vs {model_state[key].shape})")
+                
+                # Load only compatible layers
+                model.load_state_dict(filtered_state_dict, strict=False)
+
+                model.eval()
+                self.models[name] = model
+                print(f"  ✅ Successfully loaded {name}.")
+            
+            except FileNotFoundError:
+                print(f"  ❌ Checkpoint not found at: {checkpoint_path}. Skipping.")
+            except Exception as e:
+                # FIX: Detailed error reporting to show the full RuntimeError message
+                print(f"  ❌ Failed to load {name}. Error: {e.__class__.__name__}. Details: {e}. Skipping.")
+        
+        if not self.models:
+            raise RuntimeError("No models were successfully loaded. Cannot run ensemble.")
+
+    @torch.no_grad()
+    def run_ensemble(self, image_path: str) -> Dict[str, Any]:
+        """Runs inference across all loaded models and computes the ensemble prediction."""
+        
+        try:
+            image = Image.open(image_path).convert('RGB')
+            input_tensor = self.transforms(image).unsqueeze(0).to(self.device)
+        except Exception as e:
+            return {"error": f"Failed to load or process image: {e}"}
+
+        all_probs = []
+        individual_predictions = {}
+        
+        for name, model in self.models.items():
+            outputs = model(input_tensor)
+            probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
+            
+            all_probs.append(probs)
+            
+            pred_idx = np.argmax(probs)
+            pred_conf = probs[pred_idx]
+            
+            individual_predictions[name] = {
+                "class": self.class_names[pred_idx],
+                "confidence": float(pred_conf)
+            }
+
+        # Ensemble Decision (Weighted Voting)
+        # Use max confidence from each model as the weight
+        weights = np.array([np.max(probs) for probs in all_probs])
+        # Normalize weights
+        weights = weights / np.sum(weights)
+        
+        # Weighted average of probabilities
+        weighted_avg_probs = np.average(all_probs, axis=0, weights=weights)
+        ensemble_idx = np.argmax(weighted_avg_probs)
+        ensemble_confidence = weighted_avg_probs[ensemble_idx]
+        ensemble_class = self.class_names[ensemble_idx]
+
+        return {
+            "image_path": image_path,
+            "ensemble_prediction": ensemble_class,
+            "ensemble_confidence": float(ensemble_confidence),
+            "individual_predictions": individual_predictions,
+            "fracture_detected": ensemble_class != "Healthy"
+        }
+
+# ----------------------------------------------------------------------
+# --- Execution Block ---
+# ----------------------------------------------------------------------
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Multi-Model Ensemble (Cross-Validation) Agent.')
+    parser.add_argument('--image-path', required=True, help='Path to the image for inference.')
+    parser.add_argument('--checkpoints-dir', required=True, # Made required since default path was confusing
+                        help='Absolute path to the directory containing the model checkpoints (e.g., best_swin.pth).')
+    parser.add_argument('--models', type=str, default='swin,mobilenetv2,efficientnetv2,maxvit,densenet169', 
+                        help='Comma-separated names of the models to load.')
+    parser.add_argument('--num-classes', type=int, default=8)
+    parser.add_argument('--class-names', required=True, 
+                        help='Comma-separated list of class names.')
+    
+    args = parser.parse_args()
+
+    models_list = [m.strip() for m in args.models.split(',')]
+    class_names_list = [c.strip() for c in args.class_names.split(',')]
+    
+    try:
+        ensemble_agent = ModelEnsembleAgent(
+            model_names=models_list,
+            checkpoints_dir=args.checkpoints_dir,
+            num_classes=args.num_classes,
+            class_names=class_names_list
+        )
+    except RuntimeError as e:
+        print(f"\nFATAL ERROR during initialization: {e}")
+        exit(1)
+
+    result = ensemble_agent.run_ensemble(args.image_path)
+    
+    print("\n--- ENSEMBLE AGENT RESULT ---")
+    if "error" in result:
+        print(f"Error: {result['error']}")
+    else:
+        print(f"Image: {os.path.basename(result['image_path'])}")
+        print(f"FINAL ENSEMBLE PREDICTION: **{result['ensemble_prediction']}** (Confidence: {result['ensemble_confidence']:.4f})")
+        
+        print("\nIndividual Model Predictions:")
+        loaded_model_names = ensemble_agent.models.keys()
+        
+        for name in models_list:
+            if name in loaded_model_names:
+                 pred = result['individual_predictions'][name]
+                 print(f"  {name.upper():<15}: {pred['class']:<20} (Conf: {pred['confidence']:.4f})")
+            else:
+                 print(f"  {name.upper():<15}: (Skipped/Failed to Load)")
+
+    print("-----------------------------\n")

src/agents/diagnostic_agent.py

@@ -0,0 +1,142 @@
+import os
+import sys
+import argparse
+import torch
+from PIL import Image
+from typing import Dict, Any, List
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+# --- 1. CONFIGURATION ---
+from src.utils import get_device, get_model, get_transforms
+
+DEVICE = get_device()
+
+# --- 2. DIAGNOSTIC AGENT CORE ---
+
+class DiagnosticAgent:
+    def __init__(self, checkpoint_path: str, model_name: str, num_classes: int, img_size: int, class_names: List[str]):
+        self.device = DEVICE
+        self.img_size = img_size
+        self.class_names = class_names
+        self.model_name = model_name
+        
+        # 1. Load Model Architecture
+        self.model = get_model(model_name, num_classes, pretrained=False).to(self.device)
+        
+        # 2. Load Weights from Checkpoint
+        try:
+            ck = torch.load(checkpoint_path, map_location=self.device)
+            state_dict = ck.get('model_state_dict', ck)
+            self.model.load_state_dict(state_dict)
+            self.model.eval()
+            print(f"✅ Diagnostic Agent loaded model from {checkpoint_path} on {self.device}.")
+        except FileNotFoundError:
+            print(f"❌ Error: Checkpoint file not found at {checkpoint_path}")
+            exit(1)
+        except Exception as e:
+            print(f"❌ Error loading model state: {e}")
+            exit(1)
+
+        # 3. Setup Transforms
+        self.transform = get_transforms('val', self.img_size)
+
+    def run_diagnosis(self, image_path: str) -> Dict[str, Any]:
+        """
+        Runs the image classification model, detects fractures, and outputs scores.
+        
+        This method includes the fix for FileNotFoundError by resolving the path.
+        """
+        
+        # CRITICAL FIX: Convert relative path to absolute path for reliable file access
+        full_image_path = os.path.abspath(image_path)
+
+        if not os.path.exists(full_image_path):
+            # Report the original path back to the user for clarity
+            return {"error": f"Image file not found at {image_path}"}
+
+        # 1. Image Loading and Preprocessing
+        try:
+            # Use the resolved full path for PIL to open
+            img = Image.open(full_image_path).convert('RGB') 
+        except Exception as e:
+            return {"error": f"Failed to open image at {full_image_path}. Reason: {e}"}
+
+
+        img_tensor = self.transform(img).unsqueeze(0).to(self.device)
+
+        # 2. Model Inference
+        with torch.no_grad():
+            outputs = self.model(img_tensor)
+            
+        # Softmax to get probabilities (confidence scores)
+        probabilities = torch.softmax(outputs, dim=1).squeeze(0)
+        
+        # 3. Score Calculation
+        
+        predicted_idx = torch.argmax(probabilities).item()
+        confidence = probabilities[predicted_idx].item()
+        uncertainty = 1.0 - confidence
+        predicted_class_name = self.class_names[predicted_idx]
+
+        # Determine Fracture Presence (assuming 'Healthy' is a known class)
+        is_fracture_detected = (predicted_class_name != 'Healthy')
+        
+        return {
+            "image_path": image_path,
+            "fracture_detected": is_fracture_detected,
+            "predicted_class": predicted_class_name,
+            "severity_type": predicted_class_name,  # Proxy for severity
+            "confidence_score": confidence,
+            "uncertainty_score": uncertainty,
+            "all_probabilities": probabilities.cpu().numpy().tolist()
+        }
+
+# --- 3. EXECUTION ---
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Run a diagnostic agent on a single image.')
+    parser.add_argument('--image-path', type=str, required=True, help='Path to the image file to diagnose.')
+    parser.add_argument('--checkpoint', type=str, required=True, help='Path to the model checkpoint (e.g., outputs/swin_mps/best.pth)')
+    parser.add_argument('--model', type=str, default='swin', choices=['swin', 'convnext', 'densenet'])
+    parser.add_argument('--num-classes', type=int, default=8)
+    parser.add_argument('--img-size', type=int, default=224)
+    parser.add_argument('--class-names', type=str, required=True, 
+                        help='Comma-separated list of class names (e.g., "A,B,C")')
+    
+    args = parser.parse_args()
+
+    # Convert class names string to a list
+    class_names_list = [c.strip() for c in args.class_names.split(',')]
+
+    # Ensure 'Healthy' is in the list for the 'fracture_detected' check to work reliably
+    if 'Healthy' not in class_names_list:
+        print("Warning: 'Healthy' class not found in --class-names list. Fracture detection may be inaccurate.")
+    
+    # Initialize the Agent
+    agent = DiagnosticAgent(
+        checkpoint_path=args.checkpoint,
+        model_name=args.model,
+        num_classes=args.num_classes,
+        img_size=args.img_size,
+        class_names=class_names_list
+    )
+
+    # Run the Diagnosis
+    result = agent.run_diagnosis(args.image_path)
+
+    # Output Results
+    print("\n--- DIAGNOSTIC RESULTS ---")
+    if "error" in result:
+        print(f"Status: FAILED\nReason: {result['error']}")
+    else:
+        print(f"Status: SUCCESS")
+        print(f"Image: {result['image_path']}")
+        print(f"Fracture Detected: {'YES' if result['fracture_detected'] else 'NO'}")
+        print(f"Predicted Class: {result['predicted_class']}")
+        print(f"--- Scores ---")
+        print(f"Severity Type: {result['severity_type']}")
+        print(f"Confidence Score: {result['confidence_score']:.4f}")
+        print(f"Uncertainty Score: {result['uncertainty_score']:.4f}")
+    print("--------------------------\n")

src/agents/educational_agent.py

@@ -0,0 +1,148 @@
+import json
+from typing import Dict, Any
+
+class EducationalAgent:
+    """
+    Translates technical diagnosis and explanation into simple, patient-friendly terms.
+    """
+    def __init__(self, doctor_name: str = "your treating doctor"):
+        self.doctor_name = doctor_name
+        
+    def translate_to_layman_terms(self, diagnosis_result: Dict[str, Any], explanation_text: str) -> Dict[str, str]:
+        """
+        Generates simple summaries and next steps for the patient.
+
+        Args:
+            diagnosis_result: The dictionary output from DiagnosticAgent.
+            explanation_text: The string output from ExplainabilityAgent.
+
+        Returns:
+            A dictionary containing patient-friendly summary, severity, and next steps.
+        """
+        
+        # 1. Extract Key Findings
+        fracture_detected = diagnosis_result.get("fracture_detected", False)
+        predicted_class = diagnosis_result.get("predicted_class", "a specific type of injury")
+        confidence = diagnosis_result.get("confidence_score", 0.0)
+        
+        # 2. Determine Severity in Layman Terms
+        severity_map = {
+            "Healthy": "None",
+            "Greenstick": "Mild (The bone is cracked but not completely broken through.)",
+            "Transverse": "Moderate (A clean break straight across the bone.)",
+            "Oblique": "Moderate (A clean break at an angle.)",
+            "Spiral": "Serious (A twisting break that spirals around the bone.)",
+            "Comminuted": "Severe (The bone has broken into three or more pieces.)",
+            "Oblique Displaced": "Serious (The bone is broken at an angle, and the pieces are shifted out of place.)",
+            "Transverse Displaced": "Serious (The bone is broken straight across, and the pieces are shifted out of place.)",
+        }
+        
+        layman_severity = severity_map.get(predicted_class, "We need more information on this type of break.")
+        
+        # 3. Simplify the Explanation
+        
+        # Clean up the technical explanation to remove technical jargon like 'centroid' or 'activation'
+        simple_explanation = explanation_text.replace("consistent with a", "which looks like a")
+        simple_explanation = simple_explanation.replace("Confidence:", "Our computer model is highly sure (")
+        simple_explanation = simple_explanation.replace("The model's focus is", "The computer saw a clear sign of this")
+        simple_explanation = simple_explanation.replace("distal end", "end of the bone near the hand/foot")
+        simple_explanation = simple_explanation.replace("proximal end", "end of the bone near the shoulder/hip")
+        simple_explanation = simple_explanation.replace("humerus", "upper arm bone")
+        simple_explanation = simple_explanation.replace("radius", "lower arm bone")
+        simple_explanation = simple_explanation.replace("tibia", "shin bone")
+        simple_explanation = simple_explanation.replace("mild", "small")
+        simple_explanation = simple_explanation.replace("strong", "very clear")
+        
+        
+        # 4. Generate Final Summary and Next Steps
+        
+        if not fracture_detected or predicted_class == "Healthy":
+            patient_summary = (
+                f"**Great news!** Our analysis suggests your bone is **healthy** "
+                f"with high confidence ({confidence:.2f}). There are no signs of a fracture."
+            )
+            next_steps = (
+                "You can discuss your pain symptoms with your doctor, but based on this image, "
+                "a fracture is highly unlikely. No immediate orthopedic action is needed."
+            )
+        else:
+            patient_summary = (
+                f"Our computer analysis strongly indicates a **break in the bone** (a fracture). "
+                f"The specific type appears to be a **{predicted_class}** fracture."
+            )
+            
+            # Combine simple explanation and confidence
+            patient_summary += f"\n\n**What the computer saw:** {simple_explanation}"
+            patient_summary += f".\n\n**Severity Level:** {layman_severity}"
+            
+            next_steps = (
+                "This finding requires immediate medical follow-up. Please do the following:\n"
+                f"* **Do not move** the affected area.\n"
+                f"* **Immediately share these findings** with {self.doctor_name}.\n"
+                f"* Your doctor will confirm the diagnosis and determine the best treatment, "
+                "which may involve a cast, splint, or surgery."
+            )
+            
+        return {
+            "patient_summary": patient_summary,
+            "patient_severity_assessment": layman_severity,
+            "next_steps_action_plan": next_steps,
+        }
+
+# --- EXAMPLE USAGE ---
+
+if __name__ == '__main__':
+    # --- SIMULATED INPUT from Diagnostic & Explainability Agents ---
+    
+    # Example 1: Serious Fracture
+    SIMULATED_DIAGNOSIS_1 = {
+        "image_path": "fracture_image.jpg",
+        "fracture_detected": True,
+        "predicted_class": "Spiral",
+        "severity_type": "Spiral",
+        "confidence_score": 0.96,
+        "uncertainty_score": 0.04,
+        "all_probabilities": [0.01, 0.01, 0.01, 0.01, 0.01, 0.96, 0.01, 0.01]
+    }
+    SIMULATED_EXPLANATION_1 = (
+        "A fracture pattern consistent with a **Spiral** type is detected (Confidence: 0.96). "
+        "The model's focus is clear near the **middle region** of the humerus in the center. "
+        "This is based on a distinct linear focus."
+    )
+
+    # Example 2: Healthy Bone
+    SIMULATED_DIAGNOSIS_2 = {
+        "image_path": "healthy_image.jpg",
+        "fracture_detected": False,
+        "predicted_class": "Healthy",
+        "severity_type": "Healthy",
+        "confidence_score": 0.99,
+        "uncertainty_score": 0.01,
+        "all_probabilities": [0.00, 0.00, 0.99, 0.00, 0.00, 0.00, 0.00, 0.01]
+    }
+    SIMULATED_EXPLANATION_2 = (
+        "The bone appears **healthy** with high confidence (0.99). No fracture pattern was detected."
+    )
+
+
+    # --- Run Agent ---
+    
+    agent = EducationalAgent(doctor_name="Dr. Smith")
+
+    # Run Example 1
+    results_1 = agent.translate_to_layman_terms(SIMULATED_DIAGNOSIS_1, SIMULATED_EXPLANATION_1)
+    
+    print("\n--- PATIENT REPORT (FRACTURE DETECTED) ---")
+    print(f"**SUMMARY:** {results_1['patient_summary']}")
+    print("\n**ACTION PLAN:**")
+    print(results_1['next_steps_action_plan'])
+    print("-------------------------------------------\n")
+    
+    # Run Example 2
+    results_2 = agent.translate_to_layman_terms(SIMULATED_DIAGNOSIS_2, SIMULATED_EXPLANATION_2)
+
+    print("\n--- PATIENT REPORT (HEALTHY BONE) ---")
+    print(f"**SUMMARY:** {results_2['patient_summary']}")
+    print("\n**ACTION PLAN:**")
+    print(results_2['next_steps_action_plan'])
+    print("-------------------------------------\n")

src/agents/explain_agent.py

@@ -0,0 +1,164 @@
+import numpy as np
+import json
+from typing import Dict, Any
+
+# --- New Helper Function for Dynamic Testing ---
+def generate_random_heatmap(size: int = 224) -> np.ndarray:
+    """
+    Generates a randomized, plausible heatmap array for testing the agent's dynamism.
+    The heatmap will have a focused, high-intensity area somewhere random.
+    """
+    # Create a base array of zeros
+    cam_array = np.zeros((size, size), dtype=np.float32)
+    
+    # 1. Define random center and size for the activation zone
+    center_y = np.random.randint(size // 4, size * 3 // 4)
+    center_x = np.random.randint(size // 4, size * 3 // 4)
+    height = np.random.randint(30, 80)
+    width = np.random.randint(30, 80)
+    
+    # Define activation bounds (ensure they stay within the array limits)
+    y_min = max(0, center_y - height // 2)
+    y_max = min(size, center_y + height // 2)
+    x_min = max(0, center_x - width // 2)
+    x_max = min(size, center_x + width // 2)
+    
+    # 2. Apply activation with random strength
+    random_strength = np.random.uniform(0.6, 1.0)
+    cam_array[y_min:y_max, x_min:x_max] = random_strength
+    
+    # Optional: Add minor noise to make it less blocky
+    cam_array = cam_array + np.random.uniform(0, 0.1, (size, size))
+    cam_array = np.clip(cam_array, 0, 1)
+
+    return cam_array
+
+# --- Helper function for localization (No changes needed, it is dynamic) ---
+
+def calculate_heatmap_centroid(cam_array: np.ndarray, threshold: float = 0.5) -> tuple:
+    """
+    Calculates the centroid (center of mass) of the significant activation area
+    in the Grad-CAM heatmap.
+    """
+    # 1. Apply threshold to isolate the 'hot' region
+    binary_map = cam_array > threshold
+    
+    if not np.any(binary_map):
+        return (0.5, 0.5, 0.0)
+
+    # 2. Calculate coordinates and weights (activation values)
+    coords = np.argwhere(binary_map)
+    weights = cam_array[binary_map]
+
+    if len(weights) == 0:
+        return (0.5, 0.5, 0.0)
+
+    # 3. Calculate weighted average for the centroid
+    y_coords = coords[:, 0] # Rows (Y)
+    x_coords = coords[:, 1] # Columns (X)
+
+    sum_weights = np.sum(weights)
+    
+    centroid_x = np.sum(x_coords * weights) / sum_weights
+    centroid_y = np.sum(y_coords * weights) / sum_weights
+
+    # Normalize to [0, 1] based on map size
+    h, w = cam_array.shape
+    norm_x = centroid_x / w
+    norm_y = centroid_y / h
+
+    max_activation = np.max(weights)
+
+    return (norm_x, norm_y, max_activation)
+
+# --- Explainability Agent Core (No changes needed, logic is dynamic) ---
+
+class ExplainabilityAgent:
+    def __init__(self, class_names: list, body_part: str = "bone"):
+        self.class_names = class_names
+        self.body_part = body_part
+
+    def generate_explanation(self, diagnosis_result: Dict[str, Any], cam_array: np.ndarray) -> str:
+        """
+        Converts the Grad-CAM heatmap and prediction result into a textual explanation.
+        """
+        predicted_class = diagnosis_result.get("predicted_class", "Unknown")
+        confidence = diagnosis_result.get("confidence_score", 0.0)
+        
+        # 1. Analyze Heatmap
+        norm_x, norm_y, strength = calculate_heatmap_centroid(cam_array, threshold=0.4)
+        
+        # Determine general location (Simplified)
+        x_loc = "right side" if norm_x > 0.65 else ("left side" if norm_x < 0.35 else "center")
+        y_loc = "distal end" if norm_y > 0.65 else ("proximal end" if norm_y < 0.35 else "middle region")
+        
+        # 2. Build Textual Explanation based on Prediction
+        
+        if predicted_class == "Healthy":
+            if confidence > 0.90:
+                return f"The {self.body_part} appears **healthy** with high confidence ({confidence:.2f}). No fracture pattern was detected."
+            else:
+                return f"The {self.body_part} is likely **healthy** ({confidence:.2f}), though there is some low activation in the {y_loc} of the {x_loc} that warrants a closer look."
+
+        if not diagnosis_result.get("fracture_detected", True): # Default to True if key missing
+             return f"Diagnosis is **inconclusive** or data is missing."
+             
+        # 3. Explanation for Detected Fracture
+        
+        intro = f"A fracture pattern consistent with a **{predicted_class}** type is detected"
+        
+        # Strength adjective
+        if strength > 0.7:
+            strength_adj = "strong"
+        elif strength > 0.5:
+            strength_adj = "clear"
+        else:
+            strength_adj = "mild"
+            
+        # Confidence statement
+        confidence_stmt = f"(Confidence: {confidence:.2f})"
+        
+        # Location statement
+        location_stmt = f"near the **{y_loc}** of the {self.body_part} in the {x_loc}."
+        
+        # Final Assembly
+        explanation = f"{intro} {confidence_stmt}. The model's focus is {strength_adj} {location_stmt}"
+        
+        # Add a note on the type based on visual focus
+        if predicted_class in ["Transverse", "Oblique"]:
+            explanation += " This is based on a distinct linear focus."
+        
+        return explanation
+
+
+# --- 4. EXAMPLE USAGE ---
+
+if __name__ == '__main__':
+    
+    # --- SIMULATED INPUT ---
+    SIMULATED_RESULT = {
+        "image_path": "test_image.jpg",
+        "fracture_detected": True,
+        "predicted_class": "Spiral",
+        "severity_type": "Spiral",
+        "confidence_score": 0.95,
+        "uncertainty_score": 0.05,
+    }
+
+    CLASS_NAMES = ["Comminuted", "Greenstick", "Healthy", "Oblique", "Oblique Displaced", "Spiral", "Transverse", "Transverse Displaced"]
+    explainer = ExplainabilityAgent(class_names=CLASS_NAMES, body_part="humerus")
+    
+    # Run 3 times to demonstrate dynamic output
+    print("\n--- Testing Dynamic Output (Run 1: Random Heatmap) ---")
+    
+    # Use the new dynamic heatmap function!
+    dynamic_cam_1 = generate_random_heatmap()
+    explanation_text_1 = explainer.generate_explanation(SIMULATED_RESULT, dynamic_cam_1)
+    print(f"Explanation 1: {explanation_text_1}")
+
+    print("\n--- Testing Dynamic Output (Run 2: Another Random Heatmap) ---")
+    dynamic_cam_2 = generate_random_heatmap()
+    explanation_text_2 = explainer.generate_explanation(SIMULATED_RESULT, dynamic_cam_2)
+    print(f"Explanation 2: {explanation_text_2}")
+    
+    print("--------------------------------------------------\n")

src/agents/knowledge_agent.py

@@ -0,0 +1,109 @@
+import json
+from typing import Dict, Any, List
+
+# --- Pre-compiled Medical Knowledge Base (Simulated) ---
+# In a real application, this would be a large database (e.g., SQL, MongoDB, or specialized API)
+MEDICAL_KNOWLEDGE_BASE = {
+    "Comminuted": {
+        "definition": "A fracture where the bone is broken into three or more pieces.",
+        "icd_code": "S52.5",
+        "severity": "High",
+        "treatment_guidelines": ["Usually requires surgical intervention (ORIF - Open Reduction Internal Fixation).", "Long immobilization time (8-12 weeks).", "Requires physical therapy."],
+        "prognosis_notes": "Risk of non-union is higher. Full recovery may take 6+ months."
+    },
+    "Greenstick": {
+        "definition": "A partial fracture where the bone is cracked but not completely broken through. Common in children.",
+        "icd_code": "S52.3",
+        "severity": "Low-Moderate",
+        "treatment_guidelines": ["Immobilization with cast or splint.", "Careful monitoring for progression.", "Minimal surgical intervention usually needed."],
+        "prognosis_notes": "Generally good prognosis. Recovery typically within 4-6 weeks."
+    },
+    "Healthy": {
+        "definition": "No evidence of fracture. Bone appears normal.",
+        "icd_code": "Z00.0",
+        "severity": "None",
+        "treatment_guidelines": ["No treatment required.", "Continue normal activities as tolerated.", "Regular follow-up if there is persistent pain."],
+        "prognosis_notes": "Normal bone health. No intervention needed."
+    },
+    "Oblique": {
+        "definition": "A diagonal break across the bone at approximately 45 degrees.",
+        "icd_code": "S52.2",
+        "severity": "Moderate",
+        "treatment_guidelines": ["Immobilization with cast or splint.", "Regular X-rays to monitor healing.", "Physical therapy after immobilization period."],
+        "prognosis_notes": "Good prognosis with proper immobilization. Recovery typically 6-8 weeks."
+    },
+    "Oblique Displaced": {
+        "definition": "A diagonal break where the bone fragments are not aligned and have shifted out of place.",
+        "icd_code": "S52.9",
+        "severity": "Medium-High",
+        "treatment_guidelines": ["Requires reduction (closed or open).", "Often requires casting or sometimes surgery to stabilize.", "Regular X-rays to ensure proper alignment."],
+        "prognosis_notes": "Good prognosis if successfully reduced and stabilized. Recovery 8-12 weeks."
+    },
+    "Spiral": {
+        "definition": "A twisting break that spirals around the bone, typically caused by rotational forces.",
+        "icd_code": "S52.4",
+        "severity": "Serious",
+        "treatment_guidelines": ["Usually requires immobilization in a cast or brace.", "May require surgery if fragments are unstable.", "Requires extensive physical therapy."],
+        "prognosis_notes": "Variable recovery time. May take 8-16 weeks depending on severity."
+    },
+    "Transverse": {
+        "definition": "A clean break straight across the bone, perpendicular to the bone's long axis.",
+        "icd_code": "S52.1",
+        "severity": "Moderate",
+        "treatment_guidelines": ["Immobilization with cast or splint.", "Regular X-rays to monitor alignment.", "Physical therapy after healing begins."],
+        "prognosis_notes": "Good prognosis. Clean breaks typically heal well. Recovery 6-10 weeks."
+    },
+    "Transverse Displaced": {
+        "definition": "A straight break across the bone with fragments shifted out of place.",
+        "icd_code": "S52.8",
+        "severity": "Serious",
+        "treatment_guidelines": ["Requires reduction (closed or open).", "Often requires surgery to realign fragments.", "Long-term immobilization and rehabilitation."],
+        "prognosis_notes": "Good prognosis with treatment. Recovery 10-14 weeks."
+    }
+}
+
+class KnowledgeAgent:
+    def __init__(self, knowledge_base: Dict[str, Any]):
+        self.knowledge_base = knowledge_base
+
+    def get_medical_summary(self, diagnosis: str, confidence: float) -> Dict[str, Any]:
+        """
+        Retrieves and formats external medical knowledge based on the final diagnosis.
+        """
+        diagnosis = diagnosis.strip()
+
+        if diagnosis not in self.knowledge_base:
+            return {"error": "Diagnosis not found in the knowledge base."}
+        
+        # 1. Retrieve Raw Data
+        raw_data = self.knowledge_base[diagnosis]
+
+        # 2. Format Summary for Professional Use (Example output)
+        summary = {
+            "Diagnosis": diagnosis,
+            "Ensemble_Confidence": f"{confidence:.2f}",
+            "Type": raw_data.get("definition"),
+            "ICD_Code": raw_data.get("icd_code", "N/A"),
+            "Severity": raw_data.get("severity"),
+            "Guidelines": raw_data.get("treatment_guidelines")
+        }
+        
+        return summary
+
+# --- Example Usage (Integration with Cross-Validation Agent Output) ---
+if __name__ == '__main__':
+    # Assume this is the output from your cross_validation_agent:
+    cross_validation_result = {
+        "ensemble_prediction": "Oblique Displaced",
+        "ensemble_confidence": 0.85
+    }
+    
+    agent = KnowledgeAgent(MEDICAL_KNOWLEDGE_BASE)
+    
+    medical_report = agent.get_medical_summary(
+        diagnosis=cross_validation_result["ensemble_prediction"],
+        confidence=cross_validation_result["ensemble_confidence"]
+    )
+    
+    print("\n--- 🧠 KNOWLEDGE AGENT REPORT ---")
+    print(json.dumps(medical_report, indent=4))

src/analysis/__init__.py

@@ -0,0 +1,10 @@
+"""
+Analysis and visualization modules for model evaluation.
+"""
+
+# Analysis scripts can be imported individually as needed
+# from .analyze import main as analyze_results
+# from .analyze_2 import main as analyze_results_2
+# from .visualize_gradcam import main as visualize_gradcam
+
+__all__ = []

src/analysis/analyze.py

@@ -0,0 +1,104 @@
+# analyze_results.py
+import os, sys, csv, argparse, numpy as np, matplotlib.pyplot as plt
+from PIL import Image
+import torch, torch.nn as nn, torchvision.transforms as T
+import timm, torchvision.models as tvmodels
+from sklearn.metrics import precision_recall_fscore_support, confusion_matrix
+import cv2
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+from src.utils import get_device, get_model, get_transforms
+
+def load_csv(path):
+    with open(path) as f:
+        reader = csv.DictReader(f)
+        return [r for r in reader]
+
+def save_confusion(cm, labels, out_path):
+    fig, ax = plt.subplots(figsize=(8,8))
+    im = ax.imshow(cm, interpolation='nearest', cmap='Blues')
+    ax.set_xticks(range(len(labels))); ax.set_yticks(range(len(labels)))
+    ax.set_xticklabels(labels, rotation=45, ha='right'); ax.set_yticklabels(labels)
+    for i in range(len(labels)):
+        for j in range(len(labels)):
+            ax.text(j,i, str(cm[i,j]), ha='center', va='center', color='black')
+    plt.colorbar(im)
+    plt.tight_layout(); plt.savefig(out_path); plt.close(fig)
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint')
+    parser.add_argument('--test-csv')
+    parser.add_argument('--img-root', default='.')
+    parser.add_argument('--model', default='swin')
+    parser.add_argument('--img-size', default=224)
+    parser.add_argument('--class-names')
+    parser.add_argument('--out-dir', default='outputs/analysis')
+    args = parser.parse_args()
+    os.makedirs(args.out_dir, exist_ok=True)
+    class_names = [s.strip() for s in args.class_names.split(',')]
+    num_classes = len(class_names)
+    device = get_device()
+
+    model = get_model(args.model, num_classes, pretrained=False)
+    ck = torch.load(args.checkpoint, map_location='cpu')
+    model.load_state_dict(ck['model_state_dict'])
+    model.to(device); model.eval()
+
+    rows = load_csv(args.test_csv)
+    tf = get_transforms('val', args.img_size)
+    preds, trues, paths, probs = [], [], [], []
+    os.makedirs(os.path.join(args.out_dir,'examples'), exist_ok=True)
+
+    for r in rows:
+        img_path = r['image_path'] if os.path.isabs(r['image_path']) else os.path.join(args.img_root, r['image_path'])
+        img = Image.open(img_path).convert('RGB')
+        t = tf(img).unsqueeze(0).to(device)
+        with torch.no_grad():
+            out = model(t)
+            p = torch.softmax(out, dim=1).cpu().numpy()[0]
+            pred = int(p.argmax())
+        preds.append(pred); trues.append(int(r['label'])); paths.append(img_path); probs.append(p)
+
+    cm = confusion_matrix(trues, preds)
+    p, r, f1, _ = precision_recall_fscore_support(trues, preds, average=None, labels=list(range(num_classes)), zero_division=0)
+
+    # print per-class metrics
+    for i,name in enumerate(class_names):
+        print(f'{i} {name}: support={(cm[i].sum())}, prec={p[i]:.3f}, rec={r[i]:.3f}, f1={f1[i]:.3f}')
+    print('macro-f1:', np.mean(f1))
+
+    # save confusion matrix image
+    save_confusion(cm, class_names, os.path.join(args.out_dir,'confusion_matrix.png'))
+
+    # write misclassified csv
+    miscsv = os.path.join(args.out_dir,'misclassified.csv')
+    with open(miscsv,'w') as f:
+        writer = csv.writer(f); writer.writerow(['image_path','true','pred','top1','top2'])
+        for path, t, pr, prob in zip(paths,trues,preds,probs):
+            if t!=pr:
+                top2 = np.argsort(prob)[-2:][::-1].tolist()
+                writer.writerow([path, t, pr, int(np.argmax(prob)), int(top2[0])])
+
+    # Save example images for top confused pairs
+    # find the biggest off-diagonal cells
+    cm_off = cm.copy(); np.fill_diagonal(cm_off, 0)
+    flat = [(cm_off[i,j],i,j) for i in range(num_classes) for j in range(num_classes)]
+    flat = sorted(flat, reverse=True)
+    for count,i,j in flat[:6]:  # top 6 confusion pairs
+        if count==0: continue
+        pair_dir = os.path.join(args.out_dir, 'examples', f'{i}_to_{j}')
+        os.makedirs(pair_dir, exist_ok=True)
+        saved=0
+        for path,t,pred,prob in zip(paths,trues,preds,probs):
+            if t==i and pred==j and saved<10:
+                img = Image.open(path).convert('RGB')
+                img.save(os.path.join(pair_dir, os.path.basename(path)))
+                saved+=1
+
+    print('Saved misclassified list and example images in', args.out_dir)
+
+if __name__=='__main__':
+    main()

src/analysis/analyze_2.py

@@ -0,0 +1,210 @@
+import os
+import sys
+import argparse
+import time
+from pathlib import Path
+from typing import List, Dict
+
+import numpy as np
+from PIL import Image
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as T
+import torchvision.models as tvmodels
+import timm
+
+from sklearn.metrics import precision_recall_fscore_support, confusion_matrix
+import cv2
+import csv
+import matplotlib.pyplot as plt
+
+# Import necessary modules for Grad-CAM
+from pytorch_grad_cam import GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+from src.utils import get_device, get_model, get_transforms
+
+DEVICE = get_device()
+print(f"Using device: {DEVICE}")
+
+# ----------------------------- Dataset (Reusing logic from pipeline.py) -----------------------------
+
+class FractureDataset(Dataset):
+    def __init__(self, df, img_root: str = '.', transform=None):
+        self.entries = df
+        self.img_root = img_root
+        self.transform = transform
+        # CRITICAL PATH FIX: Define the redundant prefix
+        self.redundant_prefix = 'balanced_augmented_dataset/' 
+        self.redundant_prefix_len = len(self.redundant_prefix)
+
+    def __len__(self):
+        return len(len(self.entries))
+
+    def __getitem__(self, idx):
+        row = self.entries[idx]
+        img_path = row['image_path']
+        
+        # PATH CLEANING FIX: Strip the redundant prefix
+        if img_path.startswith(self.redundant_prefix):
+            img_path = img_path[self.redundant_prefix_len:]
+
+        if not os.path.isabs(img_path):
+            img_path = os.path.join(self.img_root, img_path)
+            
+        img = Image.open(img_path).convert('RGB')
+        
+        # NOTE: We return the raw image here for visualization purposes
+        raw_img = np.array(img).astype(np.float32) / 255.0
+
+        label = int(row['label']) 
+        if self.transform:
+            img = self.transform(img)
+            
+        return img, label, img_path, raw_img
+
+
+# ----------------------------- Model selection with Grad-CAM target layers -----------------------------
+
+def get_model_with_target_layer(name: str, num_classes: int, pretrained: bool=True):
+    """Get model and its target layer for Grad-CAM visualization."""
+    model = get_model(name, num_classes, pretrained=pretrained)
+    name = name.lower()
+    
+    if name.startswith('swin'):
+        # Target layer for Swin: the last layer of the last stage (blocks[-1][-1])
+        target_layer = model.layers[-1].blocks[-1].norm2
+        return model, target_layer
+    
+    if name.startswith('convnext'):
+        # Target layer for ConvNext: the last block of the feature extractor
+        target_layer = model.stages[-1]
+        return model, target_layer
+    
+    if name.startswith('densenet'):
+        # Target layer for DenseNet: features.norm5
+        target_layer = model.features.norm5
+        return model, target_layer
+
+    raise ValueError(f'Unknown target layer for model: {name}')
+
+
+# ----------------------------- Helpers: CSV loader -----------------------------
+
+def load_csv_like(path: str) -> List[Dict]:
+    rows = []
+    with open(path, 'r', encoding='utf8') as f: 
+        reader = csv.DictReader(f)
+        for r in reader:
+            rows.append(r)
+    return rows
+
+# ----------------------------- Grad-CAM Analysis -----------------------------
+
+def analyze(args):
+    device = DEVICE
+
+    # Load CSVs
+    test_rows = load_csv_like(args.test_csv)
+    
+    # Get model and the target layer for Grad-CAM
+    model, target_layer = get_model_with_target_layer(args.model, args.num_classes, pretrained=False)
+    model.to(device)
+
+    # Load checkpoint weights
+    ck = torch.load(args.checkpoint, map_location=device)
+    model.load_state_dict(ck['model_state_dict'])
+    model.eval()
+    print(f'Loaded model from {args.checkpoint} onto {device}.')
+
+    # Data setup
+    test_tf = get_transforms('val', args.img_size)
+    test_ds = FractureDataset(test_rows, img_root=args.img_root, transform=test_tf)
+    test_loader = DataLoader(test_ds, batch_size=1, shuffle=False) # Use batch size 1 for accurate CAM per image
+
+    # Initialize Grad-CAM
+    cam = GradCAM(model=model, target_layers=[target_layer], use_cuda=(device.type == 'cuda'))
+
+    # Setup output directory
+    os.makedirs(args.out_dir, exist_ok=True)
+    
+    class_names = args.class_names.split(',')
+    
+    print(f"Starting Grad-CAM analysis on {len(test_ds)} images...")
+
+    for i, (imgs, labels, img_paths, raw_imgs) in enumerate(test_loader):
+        imgs = imgs.to(device)
+        true_label = labels.item()
+        
+        # 1. Prediction and Target Setup
+        with torch.no_grad():
+            outputs = model(imgs)
+            predicted_label = outputs.softmax(dim=1).argmax(dim=1).item()
+
+        # Set the target to the PREDICTED class for visualization
+        targets = [ClassifierOutputTarget(predicted_label)]
+        
+        # 2. Generate CAM
+        grayscale_cam = cam(input_tensor=imgs, targets=targets)
+        grayscale_cam = grayscale_cam[0, :]
+        
+        # 3. Visualization
+        # raw_img is the unnormalized image [0, 1]
+        raw_img_for_viz = raw_imgs.squeeze(0).numpy()
+        visualization = show_cam_on_image(raw_img_for_viz, grayscale_cam, use_rgb=True)
+        
+        # Convert to PIL Image for saving
+        visualization_pil = Image.fromarray(cv2.cvtColor((visualization * 255).astype(np.uint8), cv2.COLOR_RGB2BGR))
+        
+        # 4. Save
+        path_obj = Path(img_paths[0])
+        class_name = class_names[true_label]
+        
+        # Define saving path
+        save_dir = os.path.join(args.out_dir, class_name)
+        os.makedirs(save_dir, exist_ok=True)
+        
+        # Determine the name with prediction/truth info
+        pred_class_name = class_names[predicted_label]
+        file_name = f'CAM_T{class_name}_P{pred_class_name}_{path_obj.name}'
+        save_path = os.path.join(save_dir, file_name)
+        
+        visualization_pil.save(save_path)
+        
+        if i % 10 == 0:
+            print(f"Processed {i+1}/{len(test_ds)}. Saved to: {save_path}")
+
+    print("Grad-CAM analysis complete. Results saved to:", args.out_dir)
+
+
+# ----------------------------- Main -----------------------------
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Run Grad-CAM analysis on test data.')
+    parser.add_argument('--checkpoint', type=str, required=True, help='Path to the model checkpoint (e.g., outputs/swin_mps/best.pth)')
+    parser.add_argument('--test-csv', type=str, required=True, help='Path to the test CSV file.')
+    parser.add_argument('--img-root', type=str, default='.', help='Root directory for images.')
+    parser.add_argument('--model', type=str, default='swin', choices=['swin','convnext'])
+    parser.add_argument('--num-classes', type=int, default=8)
+    parser.add_argument('--img-size', type=int, default=224)
+    parser.add_argument('--out-dir', type=str, default='outputs/analysis', help='Directory to save CAM visualizations.')
+    parser.add_argument('--class-names', type=str, required=True, 
+                        help='Comma-separated list of class names (e.g., "A,B,C")')
+    
+    args = parser.parse_args()
+    
+    # Check for required library dependencies
+    try:
+        import pytorch_grad_cam
+    except ImportError:
+        print("ERROR: pytorch-grad-cam library not found. Please install it:")
+        print("pip install pytorch-grad-cam")
+        exit(1)
+
+    analyze(args)

src/analysis/visualize_gradcam.py

@@ -0,0 +1,314 @@
+"""
+visualize_gradcam.py
+
+Generates Grad-CAM overlays for misclassified examples listed in a CSV (format produced earlier):
+    image_path,true,pred,top1,top2
+
+For each row this script saves a PNG with:
+  - original image
+  - Grad-CAM overlay for the **true** class
+  - Grad-CAM overlay for the **predicted** class
+  - difference overlay (pred - true)
+
+Usage:
+    python src/analysis/visualize_gradcam.py \
+      --checkpoint outputs/swin_mps/best.pth \
+      --misclassified outputs/analysis/misclassified.csv \
+      --img-root . \
+      --model swin --img-size 224 --out-dir outputs/analysis/gradcam_overlays \
+      --class-names "Comminuted,Greenstick,Healthy,Oblique,Oblique Displaced,Spiral,Transverse,Transverse Displaced"
+
+Notes:
+- Script prefers MPS (Apple Silicon) if available; if Grad-CAM backward on MPS fails it will automatically fall back to CPU for CAM computation.
+- Requires: torch, timm, torchvision, pillow, numpy, opencv-python
+
+"""
+
+import os
+import sys
+import csv
+import argparse
+from pathlib import Path
+from typing import Optional, List
+
+import numpy as np
+from PIL import Image
+import cv2
+
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+import timm
+import torchvision.models as tvmodels
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+from src.utils import get_device, get_model, get_transforms
+
+DEVICE = get_device()
+print(f"Using device: {DEVICE}")
+
+# ----------------------------- Grad-CAM Implementation -----------------------------
+
+class GradCAM:
+    """Hook-based Grad-CAM. Call with a model (in eval mode) and a target conv layer name (optional).
+    If target_layer_name is None, the last nn.Conv2d module is chosen heuristically.
+    """
+    def __init__(self, model: nn.Module, target_layer_name: Optional[str] = None):
+        self.model = model
+        self.model.eval()
+        self.activations = None
+        self.gradients = None
+        self.handles = []
+
+        # pick target layer
+        if target_layer_name is None:
+            target_layer = None
+            for n, m in reversed(list(self.model.named_modules())):
+                if isinstance(m, nn.Conv2d):
+                    target_layer_name = n
+                    target_layer = m
+                    break
+            if target_layer is None:
+                raise RuntimeError('No Conv2d layer found for Grad-CAM')
+        else:
+            target_layer = dict(self.model.named_modules()).get(target_layer_name, None)
+            if target_layer is None:
+                raise RuntimeError(f'layer name {target_layer_name} not found')
+
+        # register hooks
+        self.handles.append(target_layer.register_forward_hook(self._forward_hook))
+        # backward hook
+        try:
+            self.handles.append(target_layer.register_backward_hook(self._backward_hook))
+        except Exception:
+            # fallback for newer pytorch versions: use register_full_backward_hook if available
+            try:
+                self.handles.append(target_layer.register_full_backward_hook(self._backward_hook))
+            except Exception:
+                # some builds won't allow backward hooks; we'll compute gradients by retaining graph and reading .grad from activations
+                pass
+
+    def _forward_hook(self, module, inp, out):
+        # out: tensor shape (B,C,H,W)
+        self.activations = out.detach()
+
+    def _backward_hook(self, module, grad_in, grad_out):
+        # grad_out[0] shape (B,C,H,W)
+        self.gradients = grad_out[0].detach()
+
+    def clear(self):
+        for h in self.handles:
+            try:
+                h.remove()
+            except Exception:
+                pass
+        self.handles = []
+
+    def __call__(self, input_tensor: torch.Tensor, class_idx: Optional[int] = None, device: torch.device = torch.device('cpu')):
+        """Compute CAM for a single input tensor (1,C,H,W). Returns cam resized to input HxW in numpy [0,1]."""
+        self.model.zero_grad()
+        input_tensor = input_tensor.to(device)
+        input_tensor.requires_grad = True
+        outputs = self.model(input_tensor)
+        if class_idx is None:
+            class_idx = int(outputs.argmax(dim=1).item())
+        loss = outputs[0, class_idx]
+        loss.backward(retain_graph=True)
+
+        if self.gradients is None or self.activations is None:
+            raise RuntimeError('GradCAM failed to collect gradients/activations (hooks missing)')
+
+        grads = self.gradients[0]  # C,H,W
+        acts = self.activations[0]  # C,H,W
+        weights = grads.mean(dim=(1,2))  # C
+        cam = (weights[:, None, None] * acts).sum(dim=0).cpu().numpy()
+        cam = np.maximum(cam, 0)
+        cam = cam - cam.min()
+        if cam.max() > 0:
+            cam = cam / (cam.max() + 1e-8)
+        else:
+            cam = np.zeros_like(cam)
+        # resize to original input spatial size (assume square input)
+        H = input_tensor.shape[-2]; W = input_tensor.shape[-1]
+        cam = cv2.resize(cam, (W, H))
+        return cam
+
+
+def apply_colormap_on_image(org_img: np.ndarray, activation: np.ndarray, colormap=cv2.COLORMAP_JET, alpha=0.5):
+    """Overlay heatmap on image (org_img: HxW x 3 uint8, activation: HxW float in [0,1])"""
+    if activation is None:
+        raise ValueError('activation is None')
+    # ensure activation is 2D and in [0,1]
+    activation = np.asarray(activation)
+    if activation.ndim == 3:
+        # if somehow a channel dim exists, reduce to single channel
+        activation = activation[..., 0]
+    activation = np.clip(activation, 0.0, 1.0)
+
+    # Convert activation -> heatmap (BGR) and resize heatmap to match original image
+    heatmap = np.uint8(255 * activation)
+    heatmap = cv2.applyColorMap(heatmap, colormap)
+
+    # Resize heatmap to original image spatial size before blending
+    h, w = org_img.shape[:2]
+    heatmap = cv2.resize(heatmap, (w, h), interpolation=cv2.INTER_LINEAR)
+
+    # convert heatmap to RGB to match org_img (which is RGB)
+    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+
+    # ensure types match for addWeighted
+    org_uint8 = org_img.astype('uint8')
+    heat_uint8 = heatmap.astype('uint8')
+    overlaid = cv2.addWeighted(org_uint8, 1.0 - alpha, heat_uint8, alpha, 0)
+    return overlaid
+
+
+def pil_to_numpy(img: Image.Image):
+    arr = np.array(img.convert('RGB'))
+    return arr
+
+
+def get_transform(img_size=224):
+    return T.Compose([
+        T.Resize((img_size, img_size)),
+        T.ToTensor(),
+        T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
+    ])
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint', required=True)
+    parser.add_argument('--misclassified', required=True)
+    parser.add_argument('--img-root', default='.')
+    parser.add_argument('--model', default='swin')
+    parser.add_argument('--img-size', type=int, default=224)
+    parser.add_argument('--out-dir', default='outputs/analysis/gradcam_overlays')
+    parser.add_argument('--class-names', required=True)
+    parser.add_argument('--target-layer', default=None)
+    parser.add_argument('--max-samples', type=int, default=200, help='max misclassified rows to process')
+    args = parser.parse_args()
+
+    class_names = [c.strip() for c in args.class_names.split(',')]
+    num_classes = len(class_names)
+
+    device_pref = detect_device()
+    print('preferred device:', device_pref)
+
+    model = get_model(args.model, num_classes, pretrained=False)
+    ck = torch.load(args.checkpoint, map_location='cpu')
+    model.load_state_dict(ck['model_state_dict'])
+
+    # We'll run forward on preferred device, but if backward (for CAM) fails on MPS we'll move to CPU for CAM computation
+    model.to(device_pref)
+    model.eval()
+
+    transform = get_transform(args.img_size)
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    rows = []
+    with open(args.misclassified, 'r') as f:
+        reader = csv.DictReader(f)
+        for r in reader:
+            rows.append(r)
+    rows = rows[:args.max_samples]
+
+    # initialize GradCAM on device_pref; if backward fails, we will retry on CPU
+    gradcam = None
+    try:
+        gradcam = GradCAM(model, target_layer_name=args.target_layer)
+        cam_device = device_pref
+    except Exception as e:
+        print('GradCAM init failed on preferred device; will try CPU. Error:', e)
+        cam_device = torch.device('cpu')
+        model_cpu = get_model(args.model, num_classes, pretrained=False)
+        model_cpu.load_state_dict(ck['model_state_dict'])
+        model_cpu.to(cam_device)
+        model_cpu.eval()
+        gradcam = GradCAM(model_cpu, target_layer_name=args.target_layer)
+
+    for i, r in enumerate(rows):
+        img_path = r['image_path'] if os.path.isabs(r['image_path']) else os.path.join(args.img_root, r['image_path'])
+        true_lbl = int(r['true'])
+        pred_lbl = int(r['pred'])
+        try:
+            pil = Image.open(img_path).convert('RGB')
+        except Exception as e:
+            print('failed to open', img_path, e); continue
+
+        org_np = pil_to_numpy(pil)
+        inp = transform(pil).unsqueeze(0)
+
+        # forward on preferred device to get outputs and predicted class
+        try:
+            inp_pref = inp.to(device_pref)
+            with torch.no_grad():
+                out_pref = model(inp_pref)
+                probs = torch.softmax(out_pref, dim=1).cpu().numpy()[0]
+        except Exception as e:
+            print('forward failed on preferred device:', e)
+            # fallback to CPU forward
+            model.cpu(); inp_cpu = inp; model.eval()
+            with torch.no_grad():
+                out_cpu = model(inp_cpu)
+                probs = torch.softmax(out_cpu, dim=1).numpy()[0]
+
+        # compute CAMs on gradcam.device (cam_device)
+        cam_true = None; cam_pred = None
+        try:
+            # ensure model used for gradcam is on cam_device
+            cam_model = gradcam.model
+            # move input to cam_device
+            inp_cam = inp.to(cam_device)
+            cam_true = gradcam(inp_cam, class_idx=true_lbl, device=cam_device)
+            cam_pred = gradcam(inp_cam, class_idx=pred_lbl, device=cam_device)
+        except Exception as e:
+            print('Grad-CAM on preferred device failed for', img_path, 'error:', e)
+            # try CPU
+            try:
+                # rebuild cpu model if needed
+                cpu_dev = torch.device('cpu')
+                model_cpu = get_model(args.model, num_classes, pretrained=False)
+                model_cpu.load_state_dict(ck['model_state_dict'])
+                model_cpu.to(cpu_dev); model_cpu.eval()
+                gradcam_cpu = GradCAM(model_cpu, target_layer_name=args.target_layer)
+                cam_true = gradcam_cpu(inp.to(cpu_dev), class_idx=true_lbl, device=cpu_dev)
+                cam_pred = gradcam_cpu(inp.to(cpu_dev), class_idx=pred_lbl, device=cpu_dev)
+                gradcam_cpu.clear()
+            except Exception as e2:
+                print('Grad-CAM CPU retry failed for', img_path, e2)
+                continue
+
+        # overlay
+        try:
+            over_true = apply_colormap_on_image(org_np, cam_true, alpha=0.5)
+            over_pred = apply_colormap_on_image(org_np, cam_pred, alpha=0.5)
+            diff = cam_pred - cam_true
+            diff = (diff - diff.min()) / (diff.max() - diff.min() + 1e-8)
+            over_diff = apply_colormap_on_image(org_np, diff, alpha=0.6)
+
+            # concat: original | true | pred | diff
+            h, w, _ = org_np.shape
+            # resize overlays to original size if needed
+            over_true = cv2.resize(over_true, (w, h))
+            over_pred = cv2.resize(over_pred, (w, h))
+            over_diff = cv2.resize(over_diff, (w, h))
+            orig_bgr = cv2.cvtColor(org_np, cv2.COLOR_RGB2BGR)
+            grid = np.vstack([np.hstack([orig_bgr, cv2.cvtColor(over_true, cv2.COLOR_RGB2BGR)]),
+                              np.hstack([cv2.cvtColor(over_pred, cv2.COLOR_RGB2BGR), cv2.cvtColor(over_diff, cv2.COLOR_RGB2BGR)])])
+
+            out_name = f"{i:04d}_true{true_lbl}_pred{pred_lbl}_{os.path.basename(img_path)}.png"
+            out_path = os.path.join(args.out_dir, out_name)
+            cv2.imwrite(out_path, grid)
+        except Exception as e:
+            print('failed to create overlay for', img_path, e)
+            continue
+
+    gradcam.clear()
+    print('Saved overlays to', args.out_dir)
+
+if __name__ == '__main__':
+    main()

src/config/cloud_deployment.py

@@ -0,0 +1,253 @@
+"""
+Cloud deployment configuration for model storage and management.
+Supports AWS S3, Google Cloud Storage, and other cloud providers.
+"""
+
+import os
+import json
+from typing import Optional
+
+# ============================================================================
+# AWS S3 Configuration (if using S3 for model storage)
+# ============================================================================
+
+AWS_S3_CONFIG = {
+    "bucket": os.getenv("AWS_S3_BUCKET", "your-bucket-name"),
+    "region": os.getenv("AWS_REGION", "us-east-1"),
+    "access_key": os.getenv("AWS_ACCESS_KEY_ID", ""),
+    "secret_key": os.getenv("AWS_SECRET_ACCESS_KEY", ""),
+}
+
+# ============================================================================
+# Google Cloud Storage Configuration
+# ============================================================================
+
+GCS_CONFIG = {
+    "project_id": os.getenv("GCP_PROJECT_ID", ""),
+    "bucket": os.getenv("GCP_BUCKET", ""),
+    "credentials_json": os.getenv("GOOGLE_APPLICATION_CREDENTIALS", ""),
+}
+
+# ============================================================================
+# Model Download URLs
+# ============================================================================
+
+# These should be set as environment variables for security
+# Example for AWS S3 pre-signed URLs:
+# export SWIN_MODEL_URL="https://your-bucket.s3.amazonaws.com/best_swin.pth?..."
+
+MODEL_DOWNLOAD_URLS = {
+    "best_swin.pth": os.getenv("SWIN_MODEL_URL", ""),
+    "best_mobilenetv2.pth": os.getenv("MOBILENETV2_MODEL_URL", ""),
+    "best_densenet169.pth": os.getenv("DENSENET_MODEL_URL", ""),
+    "best_efficientnetv2.pth": os.getenv("EFFICIENTNET_MODEL_URL", ""),
+    "best_maxvit.pth": os.getenv("MAXVIT_MODEL_URL", ""),
+}
+
+# ============================================================================
+# Ollama Configuration for Cloud Deployment
+# ============================================================================
+
+OLLAMA_CONFIG = {
+    # For local deployment
+    "host": os.getenv("OLLAMA_HOST", "http://localhost:11434"),
+    "model": os.getenv("OLLAMA_MODEL", "llama3"),
+    
+    # Alternative: Use cloud-hosted LLM API instead
+    "use_cloud_api": os.getenv("USE_CLOUD_LLM", "False").lower() == "true",
+    "cloud_api_provider": os.getenv("CLOUD_LLM_PROVIDER", "openai"),  # openai, anthropic, etc
+    "cloud_api_key": os.getenv("CLOUD_LLM_API_KEY", ""),
+}
+
+# ============================================================================
+# Streamlit Cloud Configuration
+# ============================================================================
+
+STREAMLIT_CLOUD_CONFIG = {
+    "deployment_mode": os.getenv("STREAMLIT_DEPLOYMENT", "False").lower() == "true",
+    "enable_model_download": os.getenv("ENABLE_MODEL_DOWNLOAD", "True").lower() == "true",
+    "model_cache_size_mb": int(os.getenv("MODEL_CACHE_SIZE_MB", "1000")),
+}
+
+# ============================================================================
+# Helper Functions
+# ============================================================================
+
+def get_s3_client():
+    """Create AWS S3 client."""
+    try:
+        import boto3
+        return boto3.client(
+            's3',
+            region_name=AWS_S3_CONFIG["region"],
+            aws_access_key_id=AWS_S3_CONFIG["access_key"],
+            aws_secret_access_key=AWS_S3_CONFIG["secret_key"],
+        )
+    except ImportError:
+        raise ImportError("boto3 not installed. Run: pip install boto3")
+
+
+def get_gcs_client():
+    """Create Google Cloud Storage client."""
+    try:
+        from google.cloud import storage
+        return storage.Client(project=GCS_CONFIG["project_id"])
+    except ImportError:
+        raise ImportError("google-cloud-storage not installed. Run: pip install google-cloud-storage")
+
+
+def upload_models_to_s3(local_model_dir: str = "./outputs") -> dict:
+    """
+    Upload local models to AWS S3.
+    
+    Args:
+        local_model_dir: Directory containing model files
+        
+    Returns:
+        Dictionary with upload results
+    """
+    from pathlib import Path
+    
+    client = get_s3_client()
+    results = {}
+    
+    for model_file in Path(local_model_dir).glob("best_*.pth"):
+        try:
+            key = f"models/{model_file.name}"
+            print(f"Uploading {model_file.name} to S3...")
+            client.upload_file(
+                str(model_file),
+                AWS_S3_CONFIG["bucket"],
+                key,
+                Callback=None
+            )
+            results[model_file.name] = {"status": "success", "s3_key": key}
+            print(f"✅ Uploaded {model_file.name}")
+        except Exception as e:
+            results[model_file.name] = {"status": "failed", "error": str(e)}
+            print(f"❌ Failed to upload {model_file.name}: {e}")
+    
+    return results
+
+
+def upload_models_to_gcs(local_model_dir: str = "./outputs") -> dict:
+    """
+    Upload local models to Google Cloud Storage.
+    
+    Args:
+        local_model_dir: Directory containing model files
+        
+    Returns:
+        Dictionary with upload results
+    """
+    from pathlib import Path
+    
+    client = get_gcs_client()
+    bucket = client.bucket(GCS_CONFIG["bucket"])
+    results = {}
+    
+    for model_file in Path(local_model_dir).glob("best_*.pth"):
+        try:
+            blob = bucket.blob(f"models/{model_file.name}")
+            print(f"Uploading {model_file.name} to GCS...")
+            blob.upload_from_filename(str(model_file))
+            results[model_file.name] = {"status": "success", "gs_url": blob.public_url}
+            print(f"✅ Uploaded {model_file.name}")
+        except Exception as e:
+            results[model_file.name] = {"status": "failed", "error": str(e)}
+            print(f"❌ Failed to upload {model_file.name}: {e}")
+    
+    return results
+
+
+def generate_s3_presigned_urls() -> dict:
+    """Generate S3 pre-signed URLs for models."""
+    client = get_s3_client()
+    urls = {}
+    
+    for model_name in MODEL_DOWNLOAD_URLS.keys():
+        key = f"models/{model_name}"
+        try:
+            url = client.generate_presigned_url(
+                'get_object',
+                Params={'Bucket': AWS_S3_CONFIG["bucket"], 'Key': key},
+                ExpiresIn=3600 * 24 * 7  # 7 days
+            )
+            urls[model_name] = url
+        except Exception as e:
+            print(f"Error generating URL for {model_name}: {e}")
+    
+    return urls
+
+
+def print_deployment_checklist():
+    """Print deployment checklist."""
+    print("""
+╔══════════════════════════════════════════════════════════════════════════════╗
+║                    STREAMLIT CLOUD DEPLOYMENT CHECKLIST                      ║
+╚══════════════════════════════════════════════════════════════════════════════╝
+
+1. GITHUB SETUP
+   ☐ Repository pushed to GitHub
+   ☐ .gitignore excludes *.pth files
+   ☐ README.md describes the project
+   ☐ requirements-prod.txt is in root
+
+2. MODEL STORAGE (Choose one)
+   ☐ AWS S3 Setup:
+     - Created S3 bucket
+     - Uploaded models
+     - Generated pre-signed URLs
+     - Set environment variables (SWIN_MODEL_URL, etc.)
+   
+   OR
+   
+   ☐ Google Cloud Storage Setup:
+     - Created GCS bucket
+     - Uploaded models
+     - Set environment variables
+   
+   OR
+   
+   ☐ Manual Upload:
+     - Will upload models manually to Streamlit Cloud
+
+3. ENVIRONMENT VARIABLES (in Streamlit Cloud Secrets)
+   ☐ OLLAMA_HOST (if using external Ollama server)
+   ☐ OLLAMA_MODEL (default: llama3)
+   ☐ Model download URLs or credentials
+   ☐ Cloud provider credentials (if applicable)
+
+4. STREAMLIT CLOUD DEPLOYMENT
+   ☐ Created account at share.streamlit.io
+   ☐ Connected GitHub repository
+   ☐ Configured Secrets
+   ☐ Deployed app
+
+5. TESTING
+   ☐ App loads successfully
+   ☐ Models are available
+   ☐ Chat feature works (if Ollama is configured)
+   ☐ Workflow can run end-to-end
+
+═══════════════════════════════════════════════════════════════════════════════
+
+IMPORTANT NOTES:
+- Each model is ~200MB, total ~1GB
+- Streamlit Cloud max storage is ~1GB
+- Models must be downloaded/cached on startup
+- Ollama requires external server (not available in Streamlit Cloud)
+- For chat feature, consider using cloud APIs (OpenAI, Anthropic)
+
+═══════════════════════════════════════════════════════════════════════════════
+""")
+
+
+if __name__ == "__main__":
+    print("Cloud Deployment Configuration")
+    print_deployment_checklist()
+    
+    print("\n📋 Current Configuration:")
+    print(f"  Deployment Mode: {STREAMLIT_CLOUD_CONFIG['deployment_mode']}")
+    print(f"  Ollama Host: {OLLAMA_CONFIG['host']}")
+    print(f"  Use Cloud API: {OLLAMA_CONFIG['use_cloud_api']}")

src/training/__init__.py

@@ -0,0 +1,9 @@
+"""
+Training pipeline modules for model training and fine-tuning.
+"""
+
+# Training pipelines can be imported individually as needed
+# from .pipeline import main as train_pipeline
+# from .pipeline_2 import main as train_pipeline_2
+
+__all__ = []

src/training/pipeline.py

@@ -0,0 +1,394 @@
+"""
+Fracture classification pipeline — Mac MPS only with Weights & Biases logging
+
+Features:
+- Enforces MPS device on Apple Silicon (exits if not available).
+- Supports three backbones: swin, convnext, densenet (via timm / torchvision).
+- Local checkpointing (best.pth) and automatic upload of checkpoints to Weights & Biases using `wandb.save`.
+- WandB logging of train/val metrics, lr, and confusion matrix artifact.
+- Stage-2 Grad-CAM cropping and retrain supported.
+
+Usage (example):
+    python src/training/pipeline.py \
+        --train-csv data/balanced_augmented_dataset/train.csv \
+        --val-csv data/balanced_augmented_dataset/val.csv \
+        --test-csv data/balanced_augmented_dataset/test.csv \
+        --model swin --num-classes 8 --epochs 20 --batch-size 6 --img-size 224 \
+        --out-dir outputs/swin_mps --wandb-project fracture-mps --wandb-entity your_entity
+
+Notes:
+- This script *requires* MPS (Apple Silicon). It will exit if MPS is unavailable.
+- Use small batch sizes (4-8) depending on your GPU/VRAM. The Mac M4 Pro Max 36GB UM should handle moderate sizes but training is slower than CUDA GPUs.
+- For WandB: run `wandb login` beforehand or set `WANDB_API_KEY` env var.
+
+"""
+
+import os
+import sys
+import argparse
+import time
+import copy
+from pathlib import Path
+from typing import Optional, Tuple, List, Dict
+
+import numpy as np
+from PIL import Image
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as T
+import torchvision.models as tvmodels
+import timm
+
+import wandb
+from sklearn.metrics import precision_recall_fscore_support, confusion_matrix
+import cv2
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+from src.utils import require_mps, get_model, get_transforms, FractureDataset
+
+# ----------------------------- Device (MPS only) -----------------------------
+
+DEVICE = require_mps()
+print(f"Using device: {DEVICE}")
+
+# ----------------------------- Training & Evaluation -----------------------------
+
+def save_checkpoint(state, is_best, out_dir, name='checkpoint.pth', upload_to_wandb: bool=False):
+    os.makedirs(out_dir, exist_ok=True)
+    path = os.path.join(out_dir, name)
+    torch.save(state, path)
+    if is_best:
+        best_path = os.path.join(out_dir, 'best.pth')
+        torch.save(state, best_path)
+        if upload_to_wandb:
+            try:
+                wandb.save(best_path)
+                print('Uploaded best checkpoint to WandB:', best_path)
+            except Exception as e:
+                print('WandB save failed:', e)
+
+
+def train_one_epoch(model, loader, optimizer, criterion, device):
+    model.train()
+    running_loss = 0.0
+    all_preds = []
+    all_targets = []
+    for imgs, labels, _ in loader:
+        imgs = imgs.to(device)
+        labels = labels.to(device)
+        optimizer.zero_grad()
+        outputs = model(imgs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        running_loss += loss.item() * imgs.size(0)
+        preds = outputs.softmax(dim=1).argmax(dim=1)
+        all_preds.extend(preds.detach().cpu().numpy().tolist())
+        all_targets.extend(labels.detach().cpu().numpy().tolist())
+    epoch_loss = running_loss / len(loader.dataset)
+    p, r, f1, _ = precision_recall_fscore_support(all_targets, all_preds, average='macro', zero_division=0)
+    return epoch_loss, p, r, f1
+
+
+def validate(model, loader, criterion, device):
+    model.eval()
+    running_loss = 0.0
+    all_preds = []
+    all_targets = []
+    with torch.no_grad():
+        for imgs, labels, _ in loader:
+            imgs = imgs.to(device)
+            labels = labels.to(device)
+            outputs = model(imgs)
+            loss = criterion(outputs, labels)
+            running_loss += loss.item() * imgs.size(0)
+            preds = outputs.softmax(dim=1).argmax(dim=1)
+            all_preds.extend(preds.detach().cpu().numpy().tolist())
+            all_targets.extend(labels.detach().cpu().numpy().tolist())
+    epoch_loss = running_loss / len(loader.dataset)
+    p, r, f1, _ = precision_recall_fscore_support(all_targets, all_preds, average='macro', zero_division=0)
+    cm = confusion_matrix(all_targets, all_preds)
+    return epoch_loss, p, r, f1, cm
+
+# ----------------------------- Grad-CAM utilities -----------------------------
+class GradCAM:
+    def __init__(self, model: nn.Module, target_layer_name: str = None):
+        self.model = model
+        self.model.eval()
+        self.gradients = None
+        self.activations = None
+        self.hook_handles = []
+        if target_layer_name is None:
+            for n, m in reversed(list(self.model.named_modules())):
+                if isinstance(m, (nn.Conv2d,)):
+                    target_layer_name = n
+                    break
+        self.target_layer_name = target_layer_name
+        if target_layer_name is None:
+            raise ValueError('Cannot find a convolutional layer for Grad-CAM')
+        target_module = dict(self.model.named_modules())[self.target_layer_name]
+        self.hook_handles.append(target_module.register_forward_hook(self._forward_hook))
+        # Note: register_full_backward_hook not supported in all versions; use backward hook where available
+        try:
+            self.hook_handles.append(target_module.register_backward_hook(self._backward_hook))
+        except Exception:
+            pass
+
+    def _forward_hook(self, module, input, output):
+        self.activations = output.detach()
+
+    def _backward_hook(self, module, grad_in, grad_out):
+        self.gradients = grad_out[0].detach()
+
+    def __call__(self, input_tensor: torch.Tensor, class_idx: Optional[int] = None, device: torch.device = DEVICE):
+        self.model.zero_grad()
+        input_tensor = input_tensor.to(device)
+        input_tensor.requires_grad = True
+        outputs = self.model(input_tensor)
+        if class_idx is None:
+            class_idx = outputs.argmax(dim=1).item()
+        loss = outputs[0, class_idx]
+        loss.backward(retain_graph=True)
+        if self.gradients is None or self.activations is None:
+            raise RuntimeError('GradCAM failed to collect gradients/activations — try a different target layer name')
+        grads = self.gradients[0]
+        acts = self.activations[0]
+        weights = grads.mean(dim=(1,2))
+        cam = (weights[:, None, None] * acts).sum(dim=0)
+        cam = np.maximum(cam.cpu().numpy(), 0)
+        cam = cv2.resize(cam, (input_tensor.shape[-1], input_tensor.shape[-2]))
+        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
+        return cam
+
+    def close(self):
+        for h in self.hook_handles:
+            try:
+                h.remove()
+            except Exception:
+                pass
+
+# ----------------------------- Heatmap -> bbox -----------------------------
+
+def heatmap_to_bbox(cam: np.ndarray, thr: float = 0.5, min_area: int = 100):
+    H, W = cam.shape
+    thr_val = cam.max() * thr
+    mask = (cam >= thr_val).astype('uint8') * 255
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return None
+    contours = sorted(contours, key=cv2.contourArea, reverse=True)
+    for cnt in contours:
+        area = cv2.contourArea(cnt)
+        if area < min_area:
+            continue
+        x,y,w,h = cv2.boundingRect(cnt)
+        return (x, y, x+w, y+h)
+    return None
+
+# ----------------------------- Generate crops from Grad-CAM (stage 2 prep) -----------------------------
+
+def generate_crops_from_gradcam(model, entries: List[Dict], out_dir: str, transform_for_cam, device: torch.device, cam_layer: str=None, thr: float=0.5, padding: float=0.15):
+    os.makedirs(out_dir, exist_ok=True)
+    gradcam = GradCAM(model, target_layer_name=cam_layer)
+    new_entries = []
+    for i, row in enumerate(entries):
+        path = row['image_path']
+        img = Image.open(path).convert('RGB')
+        tensor = transform_for_cam(img).unsqueeze(0).to(device)
+        try:
+            cam = gradcam(tensor, class_idx=None, device=device)
+        except Exception as e:
+            print('GradCAM failed for', path, e)
+            continue
+        bbox = heatmap_to_bbox(cam, thr=thr)
+        if bbox is None:
+            w, h = img.size
+            cx, cy = w//2, h//2
+            side = int(min(w,h)*0.6)
+            xmin = max(0, cx-side//2); ymin = max(0, cy-side//2); xmax = min(w, cx+side//2); ymax = min(h, cy+side//2)
+        else:
+            xmin, ymin, xmax, ymax = bbox
+            w = xmax - xmin; h = ymax - ymin
+            px = int(w * padding); py = int(h * padding)
+            xmin = max(0, xmin - px); ymin = max(0, ymin - py); xmax = min(img.size[0], xmax + px); ymax = min(img.size[1], ymax + py)
+        crop = img.crop((xmin, ymin, xmax, ymax)).resize((224,224))
+        fname = f"crop_{i}_{os.path.basename(path)}"
+        out_path = os.path.join(out_dir, fname)
+        crop.save(out_path)
+        new_entries.append({'image_path': out_path, 'label': row['label']})
+    gradcam.close()
+    return new_entries
+
+# ----------------------------- Inference with simple TTA -----------------------------
+
+def tta_predict(model, pil_img: Image.Image, device, img_size=224):
+    base = T.Compose([T.Resize((img_size,img_size)), T.ToTensor(), T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])])
+    img1 = base(pil_img).unsqueeze(0).to(device)
+    img2 = base(pil_img.transpose(Image.FLIP_LEFT_RIGHT)).unsqueeze(0).to(device)
+    model.eval()
+    with torch.no_grad():
+        out1 = model(img1).softmax(dim=1)
+        out2 = model(img2).softmax(dim=1)
+    probs = (out1 + out2) / 2.0
+    return probs.squeeze(0).cpu().numpy()
+
+# ----------------------------- Helpers: CSV loader -----------------------------
+
+def load_csv_like(path: str) -> List[Dict]:
+    import csv
+    rows = []
+    with open(path, 'r') as f:
+        reader = csv.DictReader(f)
+        for r in reader:
+            rows.append(r)
+    return rows
+
+# ----------------------------- Main -----------------------------
+
+def main(argv=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--train-csv', type=str, help='train csv', required=True)
+    parser.add_argument('--val-csv', type=str, help='val csv', required=True)
+    parser.add_argument('--test-csv', type=str, help='test csv', required=True)
+    parser.add_argument('--img-root', type=str, default='.', help='root for images')
+    parser.add_argument('--model', type=str, default='swin', choices=['swin','convnext','densenet'])
+    parser.add_argument('--num-classes', type=int, default=8)
+    parser.add_argument('--img-size', type=int, default=224)
+    parser.add_argument('--epochs', type=int, default=20)
+    parser.add_argument('--batch-size', type=int, default=6)
+    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--weight-decay', type=float, default=1e-2)
+    parser.add_argument('--out-dir', type=str, default='outputs')
+    parser.add_argument('--checkpoint', type=str, default=None)
+    parser.add_argument('--stage2', action='store_true', help='run stage 2: generate crops from gradcam and retrain')
+    parser.add_argument('--stage2-crop-dir', type=str, default='crops')
+    parser.add_argument('--cam-layer', type=str, default=None, help='module name for Grad-CAM hook (optional)')
+
+    # wandb args
+    parser.add_argument('--wandb-project', type=str, default='fracture-mps')
+    parser.add_argument('--wandb-entity', type=str, default=None)
+    parser.add_argument('--wandb-run-name', type=str, default=None)
+    parser.add_argument('--wandb-mode', type=str, default='online', choices=['online','offline','disabled'])
+
+    args = parser.parse_args(argv)
+
+    # initialize wandb
+    if args.wandb_mode != 'disabled':
+        wandb.init(project=args.wandb_project, entity=args.wandb_entity, name=args.wandb_run_name, mode=args.wandb_mode)
+        wandb.config.update(vars(args))
+    else:
+        wandb.init(mode='disabled')
+
+    device = DEVICE
+
+    # load CSVs
+    train_rows = load_csv_like(args.train_csv)
+    val_rows = load_csv_like(args.val_csv)
+    test_rows = load_csv_like(args.test_csv)
+
+    train_tf = get_transforms('train', img_size=args.img_size)
+    val_tf = get_transforms('val', img_size=args.img_size)
+
+    model = get_model(args.model, args.num_classes, pretrained=True).to(device)
+
+    if args.checkpoint:
+        ck = torch.load(args.checkpoint, map_location=device)
+        model.load_state_dict(ck['model_state_dict'])
+        print('Loaded checkpoint', args.checkpoint)
+
+    train_ds = FractureDataset(train_rows, img_root=args.img_root, transform=train_tf)
+    val_ds = FractureDataset(val_rows, img_root=args.img_root, transform=val_tf)
+    test_ds = FractureDataset(test_rows, img_root=args.img_root, transform=val_tf)
+
+    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=False)
+    val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=False)
+    test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=False)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(1,args.epochs))
+
+    best_f1 = 0.0
+    out_dir = args.out_dir
+    os.makedirs(out_dir, exist_ok=True)
+
+    for epoch in range(args.epochs):
+        start = time.time()
+        train_loss, train_p, train_r, train_f1 = train_one_epoch(model, train_loader, optimizer, criterion, device)
+        val_loss, val_p, val_r, val_f1, cm = validate(model, val_loader, criterion, device)
+        scheduler.step()
+        is_best = val_f1 > best_f1
+        if is_best:
+            best_f1 = val_f1
+        ck_name = f'epoch_{epoch}.pth'
+        save_checkpoint({'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'val_f1': val_f1}, is_best, out_dir, name=ck_name, upload_to_wandb=(args.wandb_mode!='disabled'))
+
+        # wandb logging
+        metrics = {'epoch': epoch, 'train_loss': train_loss, 'train_macro_f1': train_f1, 'val_loss': val_loss, 'val_macro_f1': val_f1, 'lr': scheduler.get_last_lr()[0]}
+        print(f"Epoch {epoch}/{args.epochs} time={time.time()-start:.1f}s")
+        print(metrics)
+        if args.wandb_mode != 'disabled':
+            wandb.log(metrics, step=epoch)
+            # log confusion matrix as an image
+            try:
+                import matplotlib.pyplot as plt
+                fig, ax = plt.subplots(figsize=(6,6))
+                ax.imshow(cm, interpolation='nearest')
+                ax.set_title('Confusion matrix')
+                wandb.log({"confusion_matrix": wandb.Image(fig)}, step=epoch)
+                plt.close(fig)
+            except Exception as e:
+                print('Failed to log confusion matrix plot to wandb:', e)
+
+    # load best and final test evaluation
+    best_ck = os.path.join(out_dir, 'best.pth')
+    if os.path.exists(best_ck):
+        ck = torch.load(best_ck, map_location=device)
+        model.load_state_dict(ck['model_state_dict'])
+        print('Loaded best checkpoint for final evaluation')
+
+    test_loss, test_p, test_r, test_f1, test_cm = validate(model, test_loader, criterion, device)
+    print('Test results:', test_loss, test_p, test_r, test_f1)
+    np.savetxt(os.path.join(out_dir, 'confusion_matrix.txt'), test_cm, fmt='%d')
+
+    if args.wandb_mode != 'disabled':
+        # save confusion matrix as artifact
+        try:
+            wandb.log({'test_macro_f1': test_f1})
+            wandb.save(os.path.join(out_dir, 'confusion_matrix.txt'))
+        except Exception as e:
+            print('WandB final save failed:', e)
+
+    # Stage 2: Grad-CAM cropping and retrain
+    if args.stage2:
+        print('Starting Stage-2: generating crops via Grad-CAM and retraining on cropped ROIs')
+        cam_transform = T.Compose([T.Resize((224,224)), T.ToTensor(), T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])])
+        crops_out = args.stage2_crop_dir
+        new_train = generate_crops_from_gradcam(model, train_rows, out_dir=crops_out, transform_for_cam=cam_transform, device=device, cam_layer=args.cam_layer or None, thr=0.5)
+        train_ds2 = FractureDataset(new_train, transform=get_transforms('train', img_size=args.img_size))
+        train_loader2 = DataLoader(train_ds2, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=False)
+        optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(1,args.epochs//2))
+        best_f1_stage2 = 0.0
+        for epoch in range(max(5, args.epochs//2)):
+            train_loss, train_p, train_r, train_f1 = train_one_epoch(model, train_loader2, optimizer, criterion, device)
+            val_loss, val_p, val_r, val_f1, cm = validate(model, val_loader, criterion, device)
+            is_best = val_f1 > best_f1_stage2
+            if is_best:
+                best_f1_stage2 = val_f1
+            save_checkpoint({'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'val_f1': val_f1}, is_best, out_dir, name=f'stage2_epoch_{epoch}.pth', upload_to_wandb=(args.wandb_mode!='disabled'))
+            scheduler.step()
+            if args.wandb_mode != 'disabled':
+                wandb.log({'stage2_epoch': epoch, 'stage2_val_macro_f1': val_f1, 'stage2_train_macro_f1': train_f1}, step=epoch)
+        print('Stage-2 finished. Best val macro-F1:', best_f1_stage2)
+
+    print('Finished.')
+    if args.wandb_mode != 'disabled':
+        wandb.finish()
+
+if __name__ == '__main__':
+    main()

src/training/pipeline_2.py

@@ -0,0 +1,225 @@
+import os
+import sys
+import argparse
+import time
+import copy
+from pathlib import Path
+from typing import Optional, Tuple, List, Dict
+
+import numpy as np
+from PIL import Image
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as T
+import torchvision.models as tvmodels
+import timm
+
+import wandb
+from sklearn.metrics import precision_recall_fscore_support, confusion_matrix
+import cv2
+import csv 
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+from src.utils import get_device, get_model, get_transforms, FractureDataset
+
+# ----------------------------- Device Selection -----------------------------
+
+DEVICE = get_device()
+print(f"Using device: {DEVICE}")
+
+# ----------------------------- Training & Evaluation -----------------------------
+# (Omitted for brevity, but stays the same as before)
+def save_checkpoint(state, is_best, out_dir, name='checkpoint.pth', upload_to_wandb: bool=False):
+    os.makedirs(out_dir, exist_ok=True)
+    path = os.path.join(out_dir, name)
+    torch.save(state, path)
+    if is_best:
+        best_path = os.path.join(out_dir, 'best.pth')
+        torch.save(state, best_path)
+        if upload_to_wandb:
+            try:
+                wandb.save(best_path)
+                print('Uploaded best checkpoint to WandB:', best_path)
+            except Exception as e:
+                print('WandB save failed:', e)
+
+def train_one_epoch(model, loader, optimizer, criterion, device):
+    model.train()
+    running_loss = 0.0
+    all_preds = []
+    all_targets = []
+    for imgs, labels, _ in loader:
+        imgs = imgs.to(device)
+        labels = labels.to(device)
+        optimizer.zero_grad()
+        outputs = model(imgs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        running_loss += loss.item() * imgs.size(0)
+        preds = outputs.softmax(dim=1).argmax(dim=1)
+        all_preds.extend(preds.detach().cpu().numpy().tolist())
+        all_targets.extend(labels.detach().cpu().numpy().tolist())
+    epoch_loss = running_loss / len(loader.dataset)
+    p, r, f1, _ = precision_recall_fscore_support(all_targets, all_preds, average='macro', zero_division=0)
+    return epoch_loss, p, r, f1
+
+def validate(model, loader, criterion, device):
+    model.eval()
+    running_loss = 0.0
+    all_preds = []
+    all_targets = []
+    with torch.no_grad():
+        for imgs, labels, _ in loader:
+            imgs = imgs.to(device)
+            labels = labels.to(device)
+            outputs = model(imgs)
+            loss = criterion(outputs, labels)
+            running_loss += loss.item() * imgs.size(0)
+            preds = outputs.softmax(dim=1).argmax(dim=1)
+            all_preds.extend(preds.detach().cpu().numpy().tolist())
+            all_targets.extend(labels.detach().cpu().numpy().tolist())
+    epoch_loss = running_loss / len(loader.dataset)
+    p, r, f1, _ = precision_recall_fscore_support(all_targets, all_preds, average='macro', labels=list(range(outputs.shape[1])), zero_division=0)
+    cm = confusion_matrix(all_targets, all_preds, labels=list(range(outputs.shape[1])))
+    return epoch_loss, p, r, f1, cm
+
+# ----------------------------- Helpers: CSV loader -----------------------------
+# (Omitted for brevity, but stays the same as before)
+def load_csv_like(path: str) -> List[Dict]:
+    rows = []
+    with open(path, 'r', encoding='utf8') as f: 
+        reader = csv.DictReader(f)
+        for r in reader:
+            rows.append(r)
+    return rows
+
+# ----------------------------- Main -----------------------------
+
+def main(argv=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--train-csv', type=str, help='train csv', required=True)
+    parser.add_argument('--val-csv', type=str, help='val csv', required=True)
+    parser.add_argument('--test-csv', type=str, help='test csv', required=True)
+    parser.add_argument('--img-root', type=str, default='.', help='root for images')
+    parser.add_argument('--model', type=str, default='swin', choices=['swin','convnext','densenet'])
+    parser.add_argument('--num-classes', type=int, default=8)
+    parser.add_argument('--img-size', type=int, default=224)
+    parser.add_argument('--epochs', type=int, default=20)
+    parser.add_argument('--batch-size', type=int, default=6)
+    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--weight-decay', type=float, default=1e-2)
+    parser.add_argument('--out-dir', type=str, default='outputs')
+    parser.add_argument('--checkpoint', type=str, default=None)
+    parser.add_argument('--stage2', action='store_true', help='run stage 2: generate crops from gradcam and retrain')
+    parser.add_argument('--stage2-crop-dir', type=str, default='crops')
+    parser.add_argument('--cam-layer', type=str, default=None, help='module name for Grad-CAM hook (optional)')
+
+    # wandb args
+    parser.add_argument('--wandb-project', type=str, default='fracture-mps')
+    parser.add_argument('--wandb-entity', type=str, default=None)
+    parser.add_argument('--wandb-run-name', type=str, default=None)
+    parser.add_argument('--wandb-mode', type=str, default='online', choices=['online','offline','disabled'])
+
+    args = parser.parse_args(argv)
+
+    if args.wandb_mode != 'disabled':
+        wandb.init(project=args.wandb_project, entity=args.wandb_entity, name=args.wandb_run_name, mode=args.wandb_mode)
+        wandb.config.update(vars(args))
+    else:
+        wandb.init(mode='disabled') 
+
+    device = DEVICE 
+
+    train_rows = load_csv_like(args.train_csv)
+    val_rows = load_csv_like(args.val_csv)
+    test_rows = load_csv_like(args.test_csv)
+
+    train_tf = get_transforms('train', img_size=args.img_size)
+    val_tf = get_transforms('val', img_size=args.img_size)
+
+    model = get_model(args.model, args.num_classes, pretrained=True).to(device) 
+
+    if args.checkpoint:
+        ck = torch.load(args.checkpoint, map_location='cpu') 
+        state_dict = ck.get('model_state_dict', ck) 
+        model.load_state_dict(state_dict)
+        print('Loaded checkpoint', args.checkpoint)
+
+    pin_memory = device.type == 'cuda' 
+    num_workers = 0 if device.type == 'cuda' else 4 
+
+    train_ds = FractureDataset(train_rows, img_root=args.img_root, transform=train_tf)
+    val_ds = FractureDataset(val_rows, img_root=args.img_root, transform=val_tf)
+    test_ds = FractureDataset(test_rows, img_root=args.img_root, transform=val_tf)
+
+    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=num_workers, pin_memory=pin_memory)
+    val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory)
+    # FIX: Corrected typo from args.batch-size to args.batch_size
+    test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(1,args.epochs))
+
+    best_f1 = 0.0
+    out_dir = args.out_dir
+    os.makedirs(out_dir, exist_ok=True)
+
+    for epoch in range(args.epochs):
+        start = time.time()
+        train_loss, train_p, train_r, train_f1 = train_one_epoch(model, train_loader, optimizer, criterion, device)
+        val_loss, val_p, val_r, val_f1, cm = validate(model, val_loader, criterion, device)
+        scheduler.step()
+        is_best = val_f1 > best_f1
+        if is_best:
+            best_f1 = val_f1
+        ck_name = f'epoch_{epoch}.pth'
+        
+        save_checkpoint({'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'val_f1': val_f1}, is_best, out_dir, name=ck_name, upload_to_wandb=(args.wandb_mode!='disabled'))
+
+        # wandb logging
+        metrics = {'epoch': epoch, 'train_loss': train_loss, 'train_macro_f1': train_f1, 'val_loss': val_loss, 'val_macro_f1': val_f1, 'lr': scheduler.get_last_lr()[0]}
+        print(f"Epoch {epoch}/{args.epochs} time={time.time()-start:.1f}s")
+        print(metrics)
+        if args.wandb_mode != 'disabled':
+            wandb.log(metrics, step=epoch)
+            # log confusion matrix as an image
+            try:
+                import matplotlib.pyplot as plt
+                fig, ax = plt.subplots(figsize=(6,6))
+                ax.imshow(cm, interpolation='nearest')
+                ax.set_title('Confusion matrix')
+                wandb.log({"confusion_matrix": wandb.Image(fig)}, step=epoch)
+                plt.close(fig)
+            except Exception as e:
+                print('Failed to log confusion matrix plot to wandb:', e)
+
+    # load best and final test evaluation
+    best_ck = os.path.join(out_dir, 'best.pth')
+    if os.path.exists(best_ck):
+        ck = torch.load(best_ck, map_location=device)
+        model.load_state_dict(ck['model_state_dict'])
+        print('Loaded best checkpoint for final evaluation')
+
+    test_loss, test_p, test_r, test_f1, test_cm = validate(model, test_loader, criterion, device)
+    print('Test results:', test_loss, test_p, test_r, test_f1)
+    np.savetxt(os.path.join(out_dir, 'confusion_matrix.txt'), test_cm, fmt='%d')
+
+    if args.wandb_mode != 'disabled':
+        try:
+            wandb.log({'test_macro_f1': test_f1})
+            wandb.save(os.path.join(out_dir, 'confusion_matrix.txt'))
+        except Exception as e:
+            print('WandB final save failed:', e)
+
+    print('Finished.')
+    if args.wandb_mode != 'disabled':
+        wandb.finish()
+
+if __name__ == '__main__':
+    main()

src/utils/__init__.py

@@ -0,0 +1,12 @@
+from .device_utils import get_device, require_mps, DEVICE
+from .model_utils import get_model
+from .data_utils import get_transforms, FractureDataset
+
+__all__ = [
+    'get_device',
+    'require_mps',
+    'DEVICE',
+    'get_model',
+    'get_transforms',
+    'FractureDataset'
+]

src/utils/data_utils.py

@@ -0,0 +1,59 @@
+import os
+from PIL import Image
+import torch
+from torch.utils.data import Dataset
+import torchvision.transforms as T
+
+def get_transforms(split: str, img_size: int = 224):
+    """Returns train or val/test transforms."""
+    if split == 'train':
+        return T.Compose([
+            T.Resize((int(img_size*1.1), int(img_size*1.1))),
+            T.RandomResizedCrop(img_size, scale=(0.8, 1.0)),
+            T.RandomRotation(15),
+            T.RandomHorizontalFlip(),
+            T.ToTensor(),
+            T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
+        ])
+    else:
+        return T.Compose([
+            T.Resize((img_size, img_size)),
+            T.CenterCrop(img_size),
+            T.ToTensor(),
+            T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
+        ])
+
+class FractureDataset(Dataset):
+    """Dataset for fracture images with optional bounding box cropping."""
+    
+    def __init__(self, df, img_root: str = '.', transform=None, use_bbox: bool = False):
+        self.entries = df
+        self.img_root = img_root
+        self.transform = transform
+        self.use_bbox = use_bbox
+
+    def __len__(self):
+        return len(self.entries)
+
+    def __getitem__(self, idx):
+        row = self.entries[idx]
+        img_path = row['image_path']
+        
+        if not os.path.isabs(img_path):
+            img_path = os.path.join(self.img_root, img_path)
+        
+        img = Image.open(img_path).convert('RGB')
+
+        if self.use_bbox and all(k in row for k in ('bbox_xmin','bbox_ymin','bbox_xmax','bbox_ymax')):
+            xmin = int(row['bbox_xmin'])
+            ymin = int(row['bbox_ymin'])
+            xmax = int(row['bbox_xmax'])
+            ymax = int(row['bbox_ymax'])
+            img = img.crop((xmin, ymin, xmax, ymax))
+
+        label = int(row['label'])
+        
+        if self.transform:
+            img = self.transform(img)
+            
+        return img, label, img_path

src/utils/device_utils.py

@@ -0,0 +1,18 @@
+import torch
+
+def get_device():
+    """Dynamically selects CUDA, MPS, or falls back to CPU."""
+    if torch.cuda.is_available():
+        return torch.device('cuda')
+    elif getattr(torch.backends, 'mps', None) is not None and torch.backends.mps.is_available():
+        return torch.device('mps')
+    else:
+        return torch.device('cpu')
+
+def require_mps():
+    """Enforces MPS device (for Mac-only scripts)."""
+    if getattr(torch.backends, 'mps', None) is None or not torch.backends.mps.is_available():
+        raise RuntimeError('MPS (Apple Silicon) is required but not available.')
+    return torch.device('mps')
+
+DEVICE = get_device()

src/utils/model_manager.py

@@ -0,0 +1,190 @@
+"""
+Model management utility for cloud deployments.
+Handles downloading and caching models from cloud storage.
+"""
+
+import os
+import sys
+import json
+import hashlib
+from pathlib import Path
+from typing import Dict, Optional
+import requests
+
+# Add parent directory to path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+# Model registry - Update these URLs with your cloud storage
+MODEL_REGISTRY = {
+    "best_swin.pth": {
+        "size_mb": 200,
+        # Replace with your actual cloud storage URL
+        "url": os.getenv("SWIN_MODEL_URL", ""),
+        "hash": "",  # Optional: SHA256 hash for verification
+    },
+    "best_mobilenetv2.pth": {
+        "size_mb": 100,
+        "url": os.getenv("MOBILENETV2_MODEL_URL", ""),
+        "hash": "",
+    },
+    "best_densenet169.pth": {
+        "size_mb": 200,
+        "url": os.getenv("DENSENET_MODEL_URL", ""),
+        "hash": "",
+    },
+    "best_efficientnetv2.pth": {
+        "size_mb": 180,
+        "url": os.getenv("EFFICIENTNET_MODEL_URL", ""),
+        "hash": "",
+    },
+    "best_maxvit.pth": {
+        "size_mb": 220,
+        "url": os.getenv("MAXVIT_MODEL_URL", ""),
+        "hash": "",
+    },
+}
+
+MODELS_DIR = Path("./outputs")
+MODELS_DIR.mkdir(exist_ok=True)
+
+
+def check_model_exists(model_name: str) -> bool:
+    """Check if a model file exists locally."""
+    model_path = MODELS_DIR / model_name
+    return model_path.exists()
+
+
+def get_all_models_status() -> Dict[str, Dict]:
+    """Get status of all models."""
+    status = {}
+    for model_name, config in MODEL_REGISTRY.items():
+        exists = check_model_exists(model_name)
+        status[model_name] = {
+            "exists": exists,
+            "size_mb": config["size_mb"],
+            "url": config["url"],
+        }
+    return status
+
+
+def download_model(model_name: str, force: bool = False) -> bool:
+    """
+    Download a model from cloud storage.
+    
+    Args:
+        model_name: Name of the model file
+        force: Force download even if file exists
+        
+    Returns:
+        True if successful, False otherwise
+    """
+    if not force and check_model_exists(model_name):
+        print(f"✅ {model_name} already exists locally")
+        return True
+    
+    if model_name not in MODEL_REGISTRY:
+        print(f"❌ {model_name} not found in registry")
+        return False
+    
+    config = MODEL_REGISTRY[model_name]
+    url = config.get("url")
+    
+    if not url:
+        print(f"⚠️ No download URL configured for {model_name}")
+        print(f"   Set environment variable or update MODEL_REGISTRY")
+        return False
+    
+    try:
+        print(f"📥 Downloading {model_name} from cloud storage...")
+        response = requests.get(url, timeout=300, stream=True)
+        response.raise_for_status()
+        
+        model_path = MODELS_DIR / model_name
+        total_size = int(response.headers.get('content-length', 0))
+        
+        with open(model_path, 'wb') as f:
+            downloaded = 0
+            for chunk in response.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+                    downloaded += len(chunk)
+                    if total_size:
+                        percent = (downloaded / total_size) * 100
+                        print(f"  Progress: {percent:.1f}%", end='\r')
+        
+        print(f"\n✅ Successfully downloaded {model_name}")
+        return True
+    
+    except Exception as e:
+        print(f"❌ Failed to download {model_name}: {e}")
+        return False
+
+
+def download_all_models() -> Dict[str, bool]:
+    """Download all models that have URLs configured."""
+    results = {}
+    for model_name in MODEL_REGISTRY:
+        results[model_name] = download_model(model_name)
+    return results
+
+
+def initialize_models_for_deployment() -> bool:
+    """
+    Initialize models for deployment.
+    Checks if models exist, attempts download if needed.
+    
+    Returns:
+        True if all models are available, False otherwise
+    """
+    print("\n🔍 Checking model availability...")
+    status = get_all_models_status()
+    
+    all_available = True
+    for model_name, info in status.items():
+        if info["exists"]:
+            print(f"  ✅ {model_name}")
+        else:
+            print(f"  ❌ {model_name} - NOT FOUND")
+            if info["url"]:
+                print(f"     URL configured: {info['url'][:50]}...")
+            else:
+                print(f"     No download URL - configure via environment variables")
+            all_available = False
+    
+    if not all_available:
+        print("\n⚠️ Some models are missing!")
+        print("   Option 1: Configure cloud storage URLs and run: python -c 'from src.utils.model_manager import download_all_models; download_all_models()'")
+        print("   Option 2: Upload models manually to ./outputs/")
+        return False
+    
+    print("\n✅ All models are available!")
+    return True
+
+
+if __name__ == "__main__":
+    print("Model Manager - Cloud Deployment Utility")
+    print("=" * 50)
+    
+    if len(sys.argv) > 1:
+        command = sys.argv[1]
+        
+        if command == "status":
+            status = get_all_models_status()
+            print(json.dumps(status, indent=2))
+        
+        elif command == "download-all":
+            results = download_all_models()
+            print("\nDownload Results:")
+            print(json.dumps(results, indent=2))
+        
+        elif command == "check":
+            success = initialize_models_for_deployment()
+            sys.exit(0 if success else 1)
+        
+        else:
+            print(f"Unknown command: {command}")
+            print("Available commands: status, download-all, check")
+    
+    else:
+        # Default: check status
+        initialize_models_for_deployment()

src/utils/model_utils.py

@@ -0,0 +1,54 @@
+import torch.nn as nn
+import timm
+import torchvision.models as tvmodels
+
+def get_model(name: str, num_classes: int, pretrained: bool = True):
+    """Loads and adapts model architecture."""
+    name = name.lower()
+    
+    if name.startswith('swin'):
+        model = timm.create_model('swin_small_patch4_window7_224', pretrained=pretrained)
+        if hasattr(model, 'reset_classifier'):
+            model.reset_classifier(num_classes=num_classes)
+        else:
+            model.head = nn.Linear(model.head.in_features, num_classes)
+        return model
+    
+    if name.startswith('convnext'):
+        model = timm.create_model('convnext_tiny', pretrained=pretrained)
+        if hasattr(model, 'reset_classifier'):
+            model.reset_classifier(num_classes=num_classes)
+        else:
+            model.head.fc = nn.Linear(model.head.fc.in_features, num_classes)
+        return model
+    
+    if name.startswith('densenet'):
+        model = tvmodels.densenet169(pretrained=pretrained)
+        model.classifier = nn.Linear(model.classifier.in_features, num_classes)
+        return model
+    
+    if name.startswith('mobilenet'):
+        model = timm.create_model('mobilenetv2_100', pretrained=pretrained)
+        if hasattr(model, 'reset_classifier'):
+            model.reset_classifier(num_classes=num_classes)
+        else:
+            model.classifier = nn.Linear(model.classifier.in_features, num_classes)
+        return model
+    
+    if name.startswith('efficientnet'):
+        model = timm.create_model('efficientnet_b0', pretrained=pretrained)
+        if hasattr(model, 'reset_classifier'):
+            model.reset_classifier(num_classes=num_classes)
+        else:
+            model.classifier = nn.Linear(model.classifier.in_features, num_classes)
+        return model
+    
+    if name.startswith('maxvit'):
+        model = timm.create_model('maxvit_tiny_tf_224', pretrained=pretrained)
+        if hasattr(model, 'reset_classifier'):
+            model.reset_classifier(num_classes=num_classes)
+        else:
+            model.head.fc = nn.Linear(model.head.fc.in_features, num_classes)
+        return model
+    
+    raise ValueError(f'Unknown model: {name}')

streamlit_app.py

@@ -0,0 +1,66 @@
+"""
+Main entry point for Streamlit Cloud deployment.
+Streamlit Cloud looks for streamlit_app.py or app.py in the root directory.
+
+Uses the cloud-optimized version with Hugging Face Inference API.
+For local development with Ollama, use: streamlit run apps/patient_chat_app_local.py
+"""
+
+import os
+import sys
+import streamlit as st
+
+# Add src directory to Python path
+sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
+
+# Pre-initialize models check (runs once at app startup)
+@st.cache_resource
+def initialize_deployment():
+    """Initialize deployment environment and models."""
+    from src.utils.model_manager import initialize_models_for_deployment
+    
+    try:
+        models_ready = initialize_models_for_deployment()
+        return models_ready
+    except Exception as e:
+        st.error(f"Error checking models: {e}")
+        return False
+
+if __name__ == "__main__":
+    # Check model availability
+    # models_ready = initialize_deployment()
+    
+    # Import and run the cloud version with Hugging Face
+    from apps.patient_chat_app_cloud import main
+    main()
+
+# import os
+# import sys
+# import streamlit as st
+
+# # Add src directory to Python path
+# sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
+
+# # Check if we're in deployment mode
+# IS_STREAMLIT_CLOUD = os.getenv("STREAMLIT_DEPLOYMENT", "False").lower() == "true"
+
+# # Pre-initialize models check (runs once at app startup)
+# @st.cache_resource
+# def initialize_deployment():
+#     """Initialize deployment environment and models."""
+#     from src.utils.model_manager import initialize_models_for_deployment
+    
+#     try:
+#         models_ready = initialize_models_for_deployment()
+#         return models_ready
+#     except Exception as e:
+#         st.error(f"Error checking models: {e}")
+#         return False
+
+# if __name__ == "__main__":
+#     # Check model availability
+#     # models_ready = initialize_deployment()
+    
+#     # Import and run the main app
+#     from apps.patient_chat_app_local import main
+#     main()