ae OOD

app.py

@@ -1,32 +1,61 @@
-# app.py for Gradio Space with Gemini 2.5 Flash Integration + Image Samples
+# app.py for Gradio Space with Gemini 2.5 Flash Integration + OOD Detection
 import gradio as gr
 import torch
+import torch.nn as nn
 from PIL import Image
 from torchvision import transforms
 from huggingface_hub import hf_hub_download
 from model import StoolNetTriple
 import os
+import numpy as np
+import json
+import joblib
 from google import genai
 
 
+# ========================================
+# FEATURE AUTOENCODER ARCHITECTURE
+# ========================================
+class FeatureAutoencoder(nn.Module):
+    def __init__(self, input_dim: int = 512, latent_dim: int = 128, hidden_dim: int = 256):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, latent_dim),
+        )
+        self.decoder = nn.Sequential(
+            nn.Linear(latent_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, input_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        z = self.encoder(x)
+        xhat = self.decoder(z)
+        return xhat
+
+
 # ========================================
 # MODEL SETUP - Download from HF Hub
 # ========================================
-HF_TOKEN=os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN")
 
 print("Downloading model from Hugging Face Hub...")
 model_path = hf_hub_download(
     repo_id="Sid3503/stoolnet-bristol-scale-triple-attn",
     filename="pytorch_model.bin",
-    token=HF_TOKEN  # Add this parameter for authentication
+    token=HF_TOKEN
 )
 
 print("Loading model...")
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = StoolNetTriple()
-state = torch.load(model_path, map_location="cpu")
+state = torch.load(model_path, map_location=device)
 model.load_state_dict(state)
 model.eval()
-print("Model loaded successfully!")
+model = model.to(device)
+print(f"Model loaded successfully on {device}!")
 
 # Define transform
 transform = transforms.Compose([
@@ -36,12 +65,118 @@ transform = transforms.Compose([
 ])
 
 
-# Define transform
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-])
+# ========================================
+# LOAD AUTOENCODER ARTIFACTS FOR OOD DETECTION
+# ========================================
+print("Loading autoencoder artifacts for outlier detection...")
+ae_dir = 'ae_artifacts_triple_attn'
+ae_path = os.path.join(ae_dir, 'feature_ae.pth')
+scaler_path = os.path.join(ae_dir, 'scaler.pkl')
+thr_path = os.path.join(ae_dir, 'threshold.json')
+
+feature_ae = None
+feature_ae_scaler = None
+feature_ae_threshold = None
+
+try:
+    if os.path.isfile(ae_path) and os.path.isfile(scaler_path) and os.path.isfile(thr_path):
+        # Load autoencoder
+        feature_ae = FeatureAutoencoder(input_dim=512, latent_dim=128, hidden_dim=256).to(device)
+        ae_state = torch.load(ae_path, map_location=device)
+        feature_ae.load_state_dict(ae_state['model_state_dict'])
+        feature_ae.eval()
+        
+        # Load scaler
+        feature_ae_scaler = joblib.load(scaler_path)
+        
+        # Load threshold
+        with open(thr_path, 'r') as f:
+            thr_data = json.load(f)
+        feature_ae_threshold = float(thr_data.get('threshold', 0.0))
+        
+        print(f"✓ Autoencoder artifacts loaded successfully (threshold: {feature_ae_threshold:.5f})")
+    else:
+        print("⚠ Autoencoder artifact files not found. OOD detection will be disabled.")
+except Exception as e:
+    print(f"⚠ Error loading autoencoder: {e}")
+
+
+# ========================================
+# FEATURE EXTRACTION FUNCTION
+# ========================================
+@torch.no_grad()
+def extract_feature_vector(image: Image.Image) -> np.ndarray:
+    """Extract 512-D feature vector from StoolNetTriple model"""
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    
+    # Forward pass through backbone
+    x = model.conv1(image_tensor)
+    x = model.bn1(x)
+    x = model.relu(x)
+    x = model.maxpool(x)
+    x = model.layer1(x)
+    x = model.layer2(x)
+    x = model.layer3(x)
+    x = model.layer4(x)
+    
+    # Pass through triple attention
+    x, _ = model.triple_attention(x)
+    
+    # Global average pooling
+    x = model.avgpool(x)
+    feats = torch.flatten(x, 1)
+    
+    # Normalize features
+    feats = torch.nn.functional.normalize(feats, dim=1)
+    return feats.squeeze(0).cpu().numpy()
+
+
+# ========================================
+# OUTLIER DETECTION FUNCTION
+# ========================================
+def check_outlier(image: Image.Image) -> dict:
+    """Check if image is an outlier using autoencoder reconstruction error"""
+    if feature_ae is None or feature_ae_scaler is None or feature_ae_threshold is None:
+        return {
+            'is_outlier': False,
+            'ae_residual': None,
+            'ae_threshold': None,
+            'message': None
+        }
+    
+    try:
+        # Extract feature vector
+        fv = extract_feature_vector(image).astype(np.float32)
+        
+        # Standardize features
+        fv_std = feature_ae_scaler.transform(fv.reshape(1, -1)).astype(np.float32)
+        
+        # Compute reconstruction error
+        x = torch.from_numpy(fv_std).to(device)
+        with torch.no_grad():
+            recon = feature_ae(x)
+            ae_residual = torch.mean((recon - x) ** 2, dim=1).item()
+        
+        # Check if outlier
+        is_outlier = ae_residual > feature_ae_threshold
+        
+        return {
+            'is_outlier': is_outlier,
+            'ae_residual': ae_residual,
+            'ae_threshold': feature_ae_threshold,
+            'message': (
+                f"⚠️ **Outlier Detected**: This image does not appear to be a valid stool sample. "
+                f"Reconstruction error ({ae_residual:.5f}) exceeds threshold ({feature_ae_threshold:.5f})."
+            ) if is_outlier else None
+        }
+    except Exception as e:
+        print(f"Error in outlier detection: {e}")
+        return {
+            'is_outlier': False,
+            'ae_residual': None,
+            'ae_threshold': None,
+            'message': None
+        }
 
 
 # ========================================
@@ -117,6 +252,7 @@ if not GEMINI_API_KEY:
 else:
     try:
         gemini_client = genai.Client(api_key=GEMINI_API_KEY)
+        print("✓ Gemini API client initialized")
     except Exception as e:
         print(f"Error initializing Gemini client: {e}")
         gemini_client = None
@@ -133,14 +269,18 @@ def generate_gemini_prompt(type_name, shape_name, color_name, type_conf, shape_c
     color_desc = COLOR_DESCRIPTIONS[color_name]
     
     prompt = f"""You are a board-certified gastroenterologist with over 15 years of clinical experience specializing in digestive health and bowel disorders. A gastroenterologist is a medical doctor who specializes in the diagnosis and treatment of disorders of the digestive system, which includes organs from the mouth to the anus. They treat a wide range of conditions, such as irritable bowel syndrome (IBS), ulcers, heartburn, and liver diseases. To do so, they use a combination of clinical examinations, blood tests, imaging, and procedures like endoscopy.
+
 **What You Do as a Gastroenterologist:**
 **Diagnose and treat digestive disorders:** You identify and manage diseases affecting the esophagus, stomach, intestines, pancreas, liver, gallbladder, and bile ducts, with particular expertise in stool-related conditions and bowel health.
 **Perform diagnostic procedures:** You utilize tools like endoscopes, CT scans, MRIs, and ultrasounds to examine the digestive tract, and you are highly skilled in analyzing stool characteristics as diagnostic indicators.
 **Manage complex conditions:** You handle chronic diseases like hepatitis, colitis, and pancreatitis, as well as cancer and nutritional problems, with extensive experience in conditions that affect stool consistency, frequency, and appearance.
 **Provide treatment:** You prescribe medication and recommend dietary modifications, lifestyle changes, fiber supplementation, and hydration strategies specifically tailored to improve stool consistency and bowel regularity. You also perform endoscopic procedures like colonoscopy and sigmoidoscopy when stool analysis indicates the need for further investigation.
+
 You have a warm, empathetic approach and excel at explaining complex medical concepts in accessible language. Your patients appreciate your thoroughness, compassion, and practical guidance.
+
 **YOUR ROLE:**
 Provide a comprehensive, evidence-based analysis of the Bristol Stool Chart classification while maintaining a supportive, non-judgmental tone. Your goal is to educate, reassure, and empower the patient with actionable insights based on your extensive experience in stool analysis and bowel health.
+
 **REASONING FRAMEWORK (Chain of Thought):**
 Before providing your analysis, consider:
 1. What does this stool type indicate about gut transit time?
@@ -148,47 +288,68 @@ Before providing your analysis, consider:
 3. What immediate and long-term health implications should be addressed?
 4. What evidence-based interventions are most effective?
 5. When should professional medical evaluation be prioritized?
+
 ---
+
 **CURRENT CLASSIFICATION RESULT:**
+
 **Bristol Stool Type:** {type_name}
 **Description:** {type_info['description']}
 **Clinical Significance:** {type_info['condition']}
 **Confidence:** {type_conf:.1f}%
+
 **Stool Shape:** {shape_name}
 **Shape Analysis:** {shape_desc}
 **Shape Confidence:** {shape_conf:.1f}%
+
 **Stool Color:** {color_name}
 **Color Analysis:** {color_desc}
 **Color Confidence:** {color_conf:.1f}%
+
 ---
+
 **INSTRUCTIONS:**
 Provide a comprehensive gastroenterological analysis. Structure your response naturally with appropriate markdown headers (##, ###) for main sections. Do NOT include instructional labels like "Warm, Professional Greeting" or "Encouragement & Empowerment" as headers - these are content guidelines, not section titles to display.
+
 Your response should naturally flow through these topics:
+
 1. Begin with a warm, professional greeting that acknowledges the person compassionately
+
 2. Explain the Bristol Type classification in accessible terms, including what it indicates about gut transit time
+
 3. **IMPORTANT:** Integrate the shape analysis ({shape_name}: {shape_desc}) and color analysis ({color_name}: {color_desc}) into your explanation, discussing how these characteristics provide additional clinical insights
+
 4. Describe the physiological mechanisms at play (use analogies where helpful)
+
 5. Explain clinical implications, both short-term and long-term health considerations, taking into account the combination of type, shape, and color
+
 6. Provide an action plan divided into:
    - Immediate Steps (24-72 hours)
    - Medium-Term Strategy (1-4 weeks)
    - Long-term Maintenance
+
 7. Include evidence-based recommendations covering:
    - Specific dietary changes with examples
    - Hydration guidelines
    - Physical activity suggestions
    - Lifestyle modifications
    - Supplement considerations (if appropriate)
+
 8. Provide clear criteria for when to seek medical attention, including red flags (pay special attention if color is abnormal)
+
 9. Set timeline expectations - what changes to expect and when
+
 10. End with supportive closing remarks emphasizing patient agency
+
 Use descriptive section headers like "Understanding Your Classification", "What's Happening in Your Body", "Your Action Plan", etc.
+
 **TONE GUIDELINES:**
 - Warm, empathetic, non-judgmental
 - Professional yet accessible
 - Evidence-based but not overly technical
 - Empowering and action-oriented
 - Balanced (acknowledge concerns without causing alarm)
+
 **FORMATTING:**
 - Use clear markdown headers (##, ###) and subheaders
 - Bullet points for actionable items
@@ -196,6 +357,7 @@ Use descriptive section headers like "Understanding Your Classification", "What'
 - Maintain readability with short paragraphs
 - Do NOT include HTML tags - use markdown only
 - Do NOT repeat section instruction labels as headers
+
 Please provide your comprehensive analysis now."""
 
     return prompt
@@ -205,14 +367,39 @@ Please provide your comprehensive analysis now."""
 # PREDICTION FUNCTION
 # ========================================
 def predict(image):
-    """Predict Bristol type, shape, color, and get detailed AI analysis."""
+    """Predict Bristol type, shape, color, check for outliers, and get detailed AI analysis."""
     if image is None:
         return None, None, None, "⚠️ **Please upload an image to analyze.**"
     
     try:
+        # Check for outliers first
+        outlier_result = check_outlier(image)
+        
+        if outlier_result['is_outlier']:
+            # Return outlier warning without predictions
+            outlier_message = f"""## ⚠️ Invalid Image Detected
+
+{outlier_result['message']}
+
+### Quality Check Details:
+- **Reconstruction Error:** {outlier_result['ae_residual']:.5f}
+- **Threshold:** {outlier_result['ae_threshold']:.5f}
+
+**Please upload a clear image of a stool sample for accurate analysis.**
+
+---
+
+**Tips for Better Images:**
+- Ensure good lighting
+- Image should clearly show stool sample
+- Avoid blurry or unclear images
+- Image should contain actual stool content
+"""
+            return None, None, None, outlier_message
+        
         # Preprocess and run inference
         img = image.convert("RGB")
-        x = transform(img).unsqueeze(0)
+        x = transform(img).unsqueeze(0).to(device)
         
         with torch.no_grad():
             type_out, shape_out, color_out = model(x)
@@ -239,18 +426,29 @@ def predict(image):
         # Generate AI analysis
         if gemini_client is None:
             analysis = f"""## ⚠️ AI Analysis Unavailable
+
 **Classification Summary:**
+
 {BRISTOL_DESCRIPTIONS[type_name]['emoji']} **{type_name}**
 - **Description:** {BRISTOL_DESCRIPTIONS[type_name]['description']}
 - **Condition:** {BRISTOL_DESCRIPTIONS[type_name]['condition']}
 - **Confidence:** {type_conf:.1f}%
+
 **Shape:** {shape_name} ({shape_conf:.1f}% confidence)
 - {SHAPE_DESCRIPTIONS[shape_name]}
+
 **Color:** {color_name} ({color_conf:.1f}% confidence)
 - {COLOR_DESCRIPTIONS[color_name]}
+
+**Quality Check:** ✅ Image passed outlier detection
+- Reconstruction Error: {outlier_result['ae_residual']:.5f} (threshold: {outlier_result['ae_threshold']:.5f})
+
 ---
+
 *Note: Set GEMINI_API_KEY environment variable to enable detailed AI-powered gastroenterological analysis.*
+
 ### Quick Health Tips Based on Your Classification:
+
 **For {type_name}:**
 - {'Increase fiber intake and hydration' if 'Type 1' in type_name or 'Type 2' in type_name else 'Maintain balanced diet and regular exercise'}
 - {'Consider consulting a gastroenterologist if symptoms persist' if 'Type 1' in type_name or 'Type 7' in type_name else 'Monitor changes over time'}
@@ -263,29 +461,48 @@ def predict(image):
                 )
                 
                 response = gemini_client.models.generate_content(
-                    model="gemini-2.5-flash",
+                    model="gemini-2.0-flash-exp",
                     contents=prompt
                 )
                 
-                analysis = response.text
+                # Add quality check info to analysis
+                analysis = response.text + f"""
+
+---
+
+### 🔍 Quality Assurance
+
+✅ **Image Quality Check Passed**
+- **Reconstruction Error:** {outlier_result['ae_residual']:.5f}
+- **Threshold:** {outlier_result['ae_threshold']:.5f}
+- **Status:** Valid stool sample detected
+"""
                 
             except Exception as e:
                 analysis = f"""## ⚠️ Error Generating AI Analysis
+
 **Error:** {str(e)}
+
 **Classification Summary:**
+
 {BRISTOL_DESCRIPTIONS[type_name]['emoji']} **{type_name}**
 - **Description:** {BRISTOL_DESCRIPTIONS[type_name]['description']}
 - **Condition:** {BRISTOL_DESCRIPTIONS[type_name]['condition']}
 - **Confidence:** {type_conf:.1f}%
+
 **Shape:** {shape_name} ({shape_conf:.1f}% confidence)
 **Color:** {color_name} ({color_conf:.1f}% confidence)
+
+**Quality Check:** ✅ Passed (AE Error: {outlier_result['ae_residual']:.5f})
 """
         
         return type_dict, shape_dict, color_dict, analysis
         
     except Exception as e:
         error_msg = f"""## ❌ Error Processing Image
+
 **Error:** {str(e)}
+
 Please ensure you've uploaded a valid image file and try again.
 """
         return None, None, None, error_msg
@@ -321,9 +538,10 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
     gr.Markdown("""
     # 🏥 StoolNet Triple Attention — AI-Powered Bristol Stool Scale Classifier
     
-    ### Advanced Multi-Task Deep Learning Model for Comprehensive Stool Analysis
+    ### Advanced Multi-Task Deep Learning Model with Outlier Detection
     
-    This application combines **Triple Attention ResNet34** with **Google Gemini 2.5 Flash AI** for comprehensive gastroenterological assessment.
+    This application combines **Triple Attention ResNet34** with **Google Gemini 2.0 Flash AI** for comprehensive gastroenterological assessment.
+    Features built-in **outlier detection** to identify invalid or unclear images.
     """)
     
     # ========================================
@@ -342,6 +560,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
             1. Upload your image OR select from samples below
             2. Click "Analyze Sample"
             3. View results on the right
+            
+            **Note:** Images undergo automatic quality validation using autoencoder-based outlier detection.
             """)
             
             # Sample Images Section
@@ -397,9 +617,16 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
     ### 🔬 Model Information
     
     **Architecture:** Triple Attention ResNet34 (Position, Channel, Type Attention) + Multi-Task Learning  
-    **AI Analysis:** Powered by Google Gemini 2.5 Flash  
+    **AI Analysis:** Powered by Google Gemini 2.0 Flash  
+    **Outlier Detection:** Feature-space Autoencoder (trained on reference dataset)  
     **Model Repository:** [Sid3503/stoolnet-bristol-scale-triple-attn](https://huggingface.co/Sid3503/stoolnet-bristol-scale-triple-attn)
     
+    ### 🛡️ Quality Assurance Features
+    
+    - **Autoencoder-based OOD Detection**: Identifies images that don't match training distribution
+    - **Reconstruction Error Threshold**: Images with high reconstruction error are flagged as invalid
+    - **Real-time Quality Feedback**: Immediate notification if uploaded image is not a valid stool sample
+    
     ### 📚 Bristol Stool Scale Quick Reference
     
     | Type | Description | Indication |

requirements.txt

@@ -3,4 +3,5 @@ torch>=2.0.0
 torchvision>=0.15.0
 pillow>=9.0.0
 huggingface-hub>=0.19.0
-google-genai>=1.0.0
+google-genai>=1.0.0
+joblib>=1.3.0