Update app.py

app.py

@@ -2,306 +2,184 @@ import gradio as gr
 import torch
 import torch.nn as nn
 from huggingface_hub import hf_hub_download
+from transformers import CLIPModel, CLIPProcessor
 from PIL import Image
-import numpy as np
-import traceback
 import torch.nn.functional as F
-from torchvision import transforms
 
 # Your model configuration
 MODEL_REPO = "AssanaliAidarkhan/Biomedclip"
-MODEL_FILENAME = "pytorch_model.bin"
+MODEL_FILE = "pytorch_model.bin"
 
 # Global variables
 model = None
-class_labels = []
-label2id = {}
+processor = None
 id2label = {}
 
-# Simple CNN architecture (common for medical image classification)
-class SimpleCNN(nn.Module):
-    def __init__(self, num_classes):
-        super(SimpleCNN, self).__init__()
-        self.features = nn.Sequential(
-            nn.Conv2d(3, 64, kernel_size=3, padding=1),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-            nn.Conv2d(64, 128, kernel_size=3, padding=1),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-            nn.Conv2d(128, 256, kernel_size=3, padding=1),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-            nn.Conv2d(256, 512, kernel_size=3, padding=1),
-            nn.ReLU(inplace=True),
-            nn.AdaptiveAvgPool2d((7, 7))
-        )
-        self.classifier = nn.Sequential(
-            nn.Dropout(0.5),
-            nn.Linear(512 * 7 * 7, 4096),
-            nn.ReLU(inplace=True),
-            nn.Dropout(0.5),
-            nn.Linear(4096, 1000),
-            nn.ReLU(inplace=True),
-            nn.Linear(1000, num_classes)
-        )
-    
-    def forward(self, x):
-        x = self.features(x)
-        x = torch.flatten(x, 1)
-        x = self.classifier(x)
-        return x
-
-# ResNet-like architecture
-class ResNetLike(nn.Module):
-    def __init__(self, num_classes):
-        super(ResNetLike, self).__init__()
-        # Simple ResNet-like structure
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
-        self.bn1 = nn.BatchNorm2d(64)
-        self.relu = nn.ReLU(inplace=True)
-        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-        
-        # Basic blocks
-        self.layer1 = self._make_layer(64, 64, 2)
-        self.layer2 = self._make_layer(64, 128, 2, stride=2)
-        self.layer3 = self._make_layer(128, 256, 2, stride=2)
-        self.layer4 = self._make_layer(256, 512, 2, stride=2)
-        
-        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-        self.fc = nn.Linear(512, num_classes)
-    
-    def _make_layer(self, in_channels, out_channels, blocks, stride=1):
-        layers = []
-        layers.append(nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False))
-        layers.append(nn.BatchNorm2d(out_channels))
-        layers.append(nn.ReLU(inplace=True))
+class CLIPClassifier(nn.Module):
+    """Your CLIP-based classifier architecture"""
+    def __init__(self, clip_model, num_classes):
+        super(CLIPClassifier, self).__init__()
+        self.clip_model = clip_model
+        self.classifier = nn.Linear(clip_model.config.projection_dim, num_classes)
         
-        for _ in range(1, blocks):
-            layers.append(nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False))
-            layers.append(nn.BatchNorm2d(out_channels))
-            layers.append(nn.ReLU(inplace=True))
-        
-        return nn.Sequential(*layers)
-    
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.relu(x)
-        x = self.maxpool(x)
-        
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        
-        x = self.avgpool(x)
-        x = torch.flatten(x, 1)
-        x = self.fc(x)
-        return x
+    def forward(self, **inputs):
+        # Get image features from CLIP
+        outputs = self.clip_model.get_image_features(**inputs)
+        # Classify using the linear layer
+        logits = self.classifier(outputs)
+        return {'logits': logits}
 
-def try_load_with_architecture(state_dict, num_classes, architecture_name):
-    """Try loading with different architectures"""
-    
-    architectures = {
-        'simplecnn': SimpleCNN(num_classes),
-        'resnet': ResNetLike(num_classes),
-    }
-    
-    # If we know the specific architecture name, try that first
-    if architecture_name and architecture_name.lower() in architectures:
-        try:
-            model = architectures[architecture_name.lower()]
-            model.load_state_dict(state_dict)
-            return model, f"✅ Loaded with {architecture_name}"
-        except Exception as e:
-            pass
-    
-    # Try each architecture
-    for arch_name, arch_model in architectures.items():
-        try:
-            arch_model.load_state_dict(state_dict)
-            return arch_model, f"✅ Successfully loaded with {arch_name} architecture"
-        except Exception as e:
-            continue
-    
-    return None, "❌ Could not match state_dict with any architecture"
-
-def load_model():
-    """Load the BiodemCLIP model"""
-    global model, class_labels, label2id, id2label
+def load_biomedclip():
+    """Load your CLIP-based BiomedCLIP model"""
+    global model, processor, id2label
     
     try:
-        print(f"Downloading model from: {MODEL_REPO}")
+        print("🔄 Downloading model...")
         
-        # Download your model file
+        # Download model file
         model_path = hf_hub_download(
-            repo_id=MODEL_REPO,
-            filename=MODEL_FILENAME,
-            cache_dir="./model_cache"
+            repo_id=MODEL_REPO, 
+            filename=MODEL_FILE
         )
         
-        print(f"Model downloaded to: {model_path}")
+        print("✅ Model file downloaded")
         
-        # Load the saved data
-        saved_data = torch.load(model_path, map_location='cpu')
+        # Load checkpoint
+        checkpoint = torch.load(model_path, map_location='cpu')
         
-        print(f"Loaded data type: {type(saved_data)}")
+        print(f"📦 Checkpoint keys: {list(checkpoint.keys())}")
         
-        # Extract information
-        state_dict = saved_data['model_state_dict']
-        num_classes = saved_data['num_classes']
-        model_name = saved_data.get('model_name', '')
-        label2id = saved_data.get('label2id', {})
-        id2label = saved_data.get('id2label', {})
+        # Extract info
+        num_classes = checkpoint['num_classes']
+        id2label = checkpoint['id2label']
+        model_name = checkpoint.get('model_name', 'openai/clip-vit-base-patch16')
         
-        print(f"Number of classes: {num_classes}")
-        print(f"Model name: {model_name}")
-        print(f"Available labels: {list(id2label.values())[:5]}...")
+        print(f"🔢 Number of classes: {num_classes}")
+        print(f"🏷️ Classes: {list(id2label.values())}")
+        print(f"🤖 Base model: {model_name}")
         
-        # Create class labels list
-        class_labels = [id2label.get(str(i), f"Class_{i}") for i in range(num_classes)]
+        # Load CLIP processor and model
+        print("📥 Loading CLIP processor...")
+        processor = CLIPProcessor.from_pretrained(model_name)
         
-        # Try to load the model with the correct architecture
-        model, load_message = try_load_with_architecture(state_dict, num_classes, model_name)
+        print("📥 Loading CLIP model...")
+        clip_model = CLIPModel.from_pretrained(model_name)
         
-        if model is None:
-            print("❌ Failed to load model with any architecture")
-            return False
+        # Create your classifier
+        print("🔧 Creating classifier...")
+        model = CLIPClassifier(clip_model, num_classes)
         
-        print(load_message)
+        # Load your trained weights
+        print("⚙️ Loading trained weights...")
+        model.load_state_dict(checkpoint['model_state_dict'])
         
-        # Set to evaluation mode
+        # Set to eval mode
         model.eval()
         
         print("✅ Model loaded successfully!")
         return True
         
     except Exception as e:
-        print(f"Error loading model: {e}")
+        print(f"❌ Error loading model: {e}")
+        import traceback
         print(traceback.format_exc())
         return False
 
-def preprocess_image(image):
-    """Preprocess MRI image for model input"""
-    try:
-        # Convert to RGB if not already
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # Standard preprocessing pipeline
-        transform = transforms.Compose([
-            transforms.Resize((224, 224)),  # Resize to standard input size
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # ImageNet normalization
-        ])
-        
-        image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
-        
-        return image_tensor
-        
-    except Exception as e:
-        print(f"Error preprocessing image: {e}")
-        return None
-
-def classify_mri(image):
-    """Classify MRI scan"""
-    global model, class_labels, id2label
+def classify_image(image):
+    """Classify the uploaded MRI image"""
+    global model, processor, id2label
     
-    if model is None:
-        return "❌ Model not loaded! Please check the logs.", None
+    if model is None or processor is None:
+        return "❌ Model not loaded!"
     
     if image is None:
-        return "❌ Please upload an MRI scan.", None
+        return "❌ Please upload an image!"
     
     try:
-        # Preprocess the image
-        image_tensor = preprocess_image(image)
-        if image_tensor is None:
-            return "❌ Error preprocessing image.", None
+        # Convert to RGB
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Process image using CLIP processor
+        inputs = processor(images=image, return_tensors="pt")
         
         # Run inference
         with torch.no_grad():
-            outputs = model(image_tensor)
-            
-            # Apply softmax to get probabilities
-            probabilities = F.softmax(outputs, dim=-1)
-            
-            # Get top predictions
-            top_probs, top_indices = torch.topk(probabilities, k=min(5, len(class_labels)))
-            
-            # Prepare results
-            results = []
-            confidence_data = []
-            
-            for i, (prob, idx) in enumerate(zip(top_probs[0], top_indices[0])):
-                class_name = class_labels[idx.item()]
-                confidence = prob.item() * 100
-                results.append(f"{i+1}. **{class_name}**: {confidence:.2f}%")
-                confidence_data.append([class_name, confidence])
-            
-            # Get top prediction
-            top_prediction = class_labels[top_indices[0][0].item()]
-            top_confidence = top_probs[0][0].item() * 100
-            
-            # Format result text
-            result_text = f"""
-## 🔬 **MRI Classification Results**
+            outputs = model(**inputs)
+            logits = outputs['logits']
+            probabilities = F.softmax(logits, dim=1)
+        
+        # Get top predictions
+        top_probs, top_indices = torch.topk(probabilities, k=min(5, len(id2label)))
+        
+        # Format results
+        results = []
+        for i, (prob, idx) in enumerate(zip(top_probs[0], top_indices[0])):
+            class_name = id2label[str(idx.item())]  # Note: using str() for key
+            confidence = prob.item() * 100
+            results.append(f"{i+1}. **{class_name}**: {confidence:.2f}%")
+        
+        # Get top prediction
+        top_prediction = id2label[str(top_indices[0][0].item())]
+        top_confidence = top_probs[0][0].item() * 100
+        
+        result_text = f"""
+# 🔬 MRI Classification Results
 
-### 🎯 **Top Prediction:**
-**{top_prediction}** ({top_confidence:.2f}% confidence)
+## 🎯 **Top Prediction:**
+**{top_prediction}** ({top_confidence:.1f}% confidence)
 
-### 📊 **All Predictions:**
+## 📊 **All Predictions:**
 {chr(10).join(results)}
 
-### 📈 **Model Information:**
-- **Input Size:** 224×224 pixels
-- **Model:** BiomedCLIP 
-- **Classes:** {len(class_labels)} categories
-- **Available Classes:** {', '.join(class_labels)}
+## 📈 **Model Info:**
+- **Architecture:** CLIP-based classifier
+- **Classes:** {len(id2label)} categories  
+- **Input processed:** ✅
 
-### 💡 **Interpretation:**
-{get_interpretation(top_prediction, top_confidence)}
-            """
-            
-            return result_text, confidence_data
-                
+## 💡 **Confidence Level:**
+{get_confidence_interpretation(top_confidence)}
+        """
+        
+        return result_text
+        
     except Exception as e:
-        error_msg = f"❌ **Error during classification:** {str(e)}\n\n{traceback.format_exc()}"
-        return error_msg, None
-
-def get_interpretation(prediction, confidence):
-    """Provide interpretation based on prediction"""
-    if confidence >= 90:
-        return f"🟢 **High Confidence**: The model is very confident this is {prediction.lower()}."
-    elif confidence >= 70:
-        return f"🟡 **Good Confidence**: The model believes this is likely {prediction.lower()}."
-    elif confidence >= 50:
-        return f"🟠 **Moderate Confidence**: The model suggests this might be {prediction.lower()}, but consider additional analysis."
+        return f"❌ Classification error: {str(e)}"
+
+def get_confidence_interpretation(confidence):
+    """Interpret confidence level"""
+    if confidence >= 80:
+        return "🟢 **High confidence** - Strong classification result"
+    elif confidence >= 60:
+        return "🟡 **Good confidence** - Reliable result"  
+    elif confidence >= 40:
+        return "🟠 **Moderate confidence** - Consider additional analysis"
     else:
-        return f"🔴 **Low Confidence**: The model is uncertain. Manual review recommended."
+        return "🔴 **Low confidence** - Uncertain result, manual review recommended"
 
-# Load model on startup
-print("Initializing BiomedCLIP model...")
-model_loaded = load_model()
+# Load model at startup
+print("🚀 Loading BiomedCLIP model...")
+model_loaded = load_biomedclip()
 
 # Create Gradio interface
-with gr.Blocks(title="MRI Classification with BiomedCLIP", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="BiomedCLIP MRI Classifier") as app:
+    
     gr.Markdown("""
-    # 🧠 MRI Classification with BiomedCLIP
+    # 🧠 BiomedCLIP MRI Classifier
     
-    Upload an MRI scan to get automated classification results.
+    Upload an MRI scan for automated medical image classification.
     
-    **Model:** AssanaliAidarkhan/Biomedclip
+    **Model:** AssanaliAidarkhan/Biomedclip (CLIP-based)
     """)
     
-    if not model_loaded:
-        gr.Markdown("⚠️ **Warning: Model failed to load. Check the logs for details.**")
-    else:
+    if model_loaded:
         gr.Markdown("✅ **Model loaded successfully!**")
+    else:
+        gr.Markdown("❌ **Model failed to load - check logs below**")
     
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column():
+            # This is where you upload your image! 👇
             image_input = gr.Image(
                 type="pil", 
                 label="📸 Upload MRI Scan",
@@ -309,58 +187,32 @@ with gr.Blocks(title="MRI Classification with BiomedCLIP", theme=gr.themes.Soft(
             )
             
             classify_btn = gr.Button("🔍 Classify MRI", variant="primary", size="lg")
-            clear_btn = gr.Button("🗑️ Clear", variant="secondary")
+            clear_btn = gr.Button("🗑️ Clear")
         
-        with gr.Column(scale=1):
-            result_output = gr.Markdown(label="📊 Classification Results")
-            
-            # Confidence chart
-            confidence_plot = gr.BarPlot(
-                x="class",
-                y="confidence", 
-                title="Confidence Scores by Class",
-                x_title="Medical Condition",
-                y_title="Confidence (%)",
-                width=500,
-                height=300
-            )
-    
-    # Event handlers
-    def classify_and_plot(image):
-        text_result, plot_data = classify_mri(image)
-        if plot_data:
-            plot_df = [{"class": item[0], "confidence": item[1]} for item in plot_data]
-            return text_result, plot_df
-        return text_result, None
+        with gr.Column():
+            result_output = gr.Markdown(label="📊 Results")
     
+    # Button actions
     classify_btn.click(
-        fn=classify_and_plot,
+        fn=classify_image,
         inputs=image_input,
-        outputs=[result_output, confidence_plot]
+        outputs=result_output
     )
     
     clear_btn.click(
-        fn=lambda: [None, "", None],
-        inputs=[],
-        outputs=[image_input, result_output, confidence_plot]
+        fn=lambda: [None, ""],
+        outputs=[image_input, result_output]
     )
     
-    # Instructions
     gr.Markdown("""
-    ### 📋 Instructions:
-    1. **Upload an MRI scan image** (JPEG, PNG, etc.) using the image upload area above ☝️
-    2. Click "🔍 Classify MRI" to get results
-    3. View the classification results and confidence scores
-    
-    ### 🏥 Model Information:
-    - Automatically detects the number of classes from your trained model
-    - Uses the exact class labels from your training data
-    - Applies standard medical image preprocessing
+    ### 📋 How to Use:
+    1. **Click the image area above** or **drag & drop** your MRI image
+    2. Click "🔍 Classify MRI" 
+    3. View results below
     
-    ### ⚠️ Medical Disclaimer:
-    This tool is for research purposes only and should not replace professional medical diagnosis.
-    Always consult qualified medical professionals for clinical decisions.
+    ### 🏥 Medical Disclaimer:
+    For research purposes only. Not for clinical diagnosis.
     """)
 
 if __name__ == "__main__":
-    demo.launch()
+    app.launch()