Update app.py

app.py

@@ -1,55 +1,22 @@
 
 
 import torch
-from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration, BitsAndBytesConfig
+from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
 import gradio as gr
 from PIL import Image
 import re
-import os
 from typing import List, Tuple
 
-# Configuration for 4-bit quantization (if GPU available)
-quant_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True
-)
-
 class RiverPollutionAnalyzer:
     def __init__(self):
-        try:
-            # Initialize model with fallback for CPU
-            self.processor = InstructBlipProcessor.from_pretrained(
-                "Salesforce/instructblip-flan-t5-xl",
-                cache_dir="model_cache"
-            )
-            
-            if torch.cuda.is_available():
-                self.model = InstructBlipForConditionalGeneration.from_pretrained(
-                    "Salesforce/instructblip-flan-t5-xl",
-                    device_map="auto",
-                    quantization_config=quant_config,
-                    torch_dtype=torch.float16,
-                    cache_dir="model_cache"
-                )
-                self.device = "cuda"
-                self.status = "✅ Model loaded (4-bit GPU)"
-            else:
-                self.model = InstructBlipForConditionalGeneration.from_pretrained(
-                    "Salesforce/instructblip-flan-t5-xl",
-                    device_map="auto",
-                    torch_dtype=torch.float32,
-                    cache_dir="model_cache",
-                    low_cpu_mem_usage=True
-                )
-                self.device = "cpu"
-                self.status = "⚠️ Model loaded (CPU mode - slower)"
-                
-        except Exception as e:
-            self.model = None
-            self.status = f"❌ Model loading failed: {str(e)}"
-            print(self.status)
+        # Initialize model with 4-bit quantization
+        self.processor = InstructBlipProcessor.from_pretrained("Salesforce/instructblip-vicuna-7b")
+        self.model = InstructBlipForConditionalGeneration.from_pretrained(
+            "Salesforce/instructblip-vicuna-7b",
+            device_map="auto",
+            torch_dtype=torch.float16,
+            load_in_4bit=True
+        )
 
         self.pollutants = [
             "plastic waste", "chemical foam", "industrial discharge",
@@ -72,15 +39,9 @@ class RiverPollutionAnalyzer:
         }
 
     def analyze_image(self, image):
-        """Analyze river pollution with device-aware processing"""
-        if not self.model:
-            return "Model not loaded. Please check logs."
-            
+        """Analyze river pollution with robust parsing"""
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image)
-        
-        # Resize for efficiency
-        image = image.resize((512, 512))
 
         prompt = """Analyze this river pollution scene and provide:
 1. List ALL visible pollutants ONLY from: [plastic waste, chemical foam, industrial discharge, sewage water, oil spill, organic debris, construction waste, medical waste, floating trash, algal bloom, toxic sludge, agricultural runoff]
@@ -90,31 +51,27 @@ Respond EXACTLY in this format:
 Pollutants: [comma separated list]
 Severity: [number]"""
 
-        try:
-            inputs = self.processor(
-                images=image,
-                text=prompt,
-                return_tensors="pt"
-            ).to(self.model.device)
-
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=150,  # Reduced for stability
-                    temperature=0.5,
-                    top_p=0.85,
-                    do_sample=True
-                )
-
-            analysis = self.processor.batch_decode(outputs, skip_special_tokens=True)[0]
-            pollutants, severity = self._parse_response(analysis)
-            return self._format_analysis(pollutants, severity)
-        except Exception as e:
-            return f"⚠️ Analysis error: {str(e)}"
-
-    # [Keep all existing helper methods unchanged]
+        inputs = self.processor(
+            images=image,
+            text=prompt,
+            return_tensors="pt"
+        ).to("cuda", torch.float16)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=200,
+                temperature=0.5,
+                top_p=0.85,
+                do_sample=True
+            )
+
+        analysis = self.processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        pollutants, severity = self._parse_response(analysis)
+        return self._format_analysis(pollutants, severity)
+
     def _parse_response(self, analysis: str) -> Tuple[List[str], int]:
-        """Same parsing logic as before"""
+        """Robust parsing of model response"""
         pollutants = []
         severity = 3
 
@@ -123,6 +80,7 @@ Severity: [number]"""
             r'(?i)(pollutants?|contaminants?)[:\s]*\[?(.*?)(?:\]|Severity|severity|$)',
             analysis
         )
+
         if pollutant_match:
             pollutants_str = pollutant_match.group(2).strip()
             pollutants = [
@@ -132,7 +90,11 @@ Severity: [number]"""
             ]
 
         # Extract severity
-        severity_match = re.search(r'(?i)(severity|level)[:\s]*(\d{1,2})', analysis)
+        severity_match = re.search(
+            r'(?i)(severity|level)[:\s]*(\d{1,2})',
+            analysis
+        )
+
         if severity_match:
             try:
                 severity = min(max(int(severity_match.group(2)), 1), 10)
@@ -144,7 +106,7 @@ Severity: [number]"""
         return pollutants, severity
 
     def _calculate_severity(self, pollutants: List[str]) -> int:
-        """Same severity calculation"""
+        """Weighted severity calculation"""
         if not pollutants:
             return 1
 
@@ -159,7 +121,7 @@ Severity: [number]"""
         return min(10, max(1, round(avg_weight * 3)))
 
     def _format_analysis(self, pollutants: List[str], severity: int) -> str:
-        """Same formatting"""
+        """Generate formatted report"""
         severity_bar = f"""📊 Severity: {severity}/10
 {"█" * severity}{"░" * (10 - severity)}
 {self.severity_descriptions.get(severity, '')}"""
@@ -171,57 +133,21 @@ Severity: [number]"""
 {pollutants_list}
 {severity_bar}"""
 
-    def analyze_chat(self, message: str) -> str:
-        """Handle chat questions"""
-        if any(word in message.lower() for word in ["hello", "hi", "hey"]):
-            return "Hello! I'm a river pollution analyzer. Ask me about pollution types."
-        elif "pollution" in message.lower():
-            return "Common river pollutants: plastic waste, chemical foam, industrial discharge, sewage water, oil spills."
-        else:
-            return "I can answer questions about river pollution. Try asking about pollution types."
-
 # Initialize analyzer
 analyzer = RiverPollutionAnalyzer()
 
-# Gradio Interface
+import gradio as gr
+
+ # Import your actual analyzer
+
 css = """
-.header {
-    text-align: center;
-    padding: 20px;
-    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
-    border-radius: 10px;
-    margin-bottom: 20px;
-}
-.side-by-side {
-    display: flex;
-    gap: 20px;
-}
-.left-panel, .right-panel {
-    flex: 1;
-}
-.analysis-box {
-    padding: 20px;
-    background: #f8f9fa;
-    border-radius: 10px;
-    margin-top: 20px;
-    border: 1px solid #dee2e6;
-}
-.chat-container {
-    background: #f8f9fa;
-    padding: 20px;
-    border-radius: 10px;
-    height: 100%;
-}
-.dark .analysis-box, .dark .chat-container {
-    background: #2a2a2a;
-    border-color: #444;
-}
+/* (Keep all your CSS styles) */
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     with gr.Column(elem_classes="header"):
         gr.Markdown("# 🌍 River Pollution Analyzer")
-        gr.Markdown(f"### {analyzer.status}")
+        gr.Markdown("### AI-powered water pollution detection")
 
     with gr.Row(elem_classes="side-by-side"):
         # Left Panel
@@ -231,53 +157,54 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 analyze_btn = gr.Button("🔍 Analyze Pollution", variant="primary")
             
             with gr.Group(elem_classes="analysis-box"):
-                gr.Markdown("### 📊 Analysis Report")
+                gr.Markdown("### 📊 Analysis report")
                 analysis_output = gr.Markdown()
 
         # Right Panel
         with gr.Column(elem_classes="right-panel"):
             with gr.Group(elem_classes="chat-container"):
-                chatbot = gr.Chatbot(label="Pollution Q&A", height=400)
+                chatbot = gr.Chatbot(label="Pollution Analysis Q&A", height=400)
                 with gr.Row():
-                    chat_input = gr.Textbox(
-                        placeholder="Ask about pollution sources...",
-                        label="Your Question",
-                        container=False,
-                        scale=5
-                    )
+                    chat_input = gr.Textbox(placeholder="Ask about pollution sources...", 
+                                         label="Your Question", container=False, scale=5)
                     chat_btn = gr.Button("💬 Ask", variant="secondary", scale=1)
-                clear_btn = gr.Button("🧹 Clear Chat", size="sm")
+                clear_btn = gr.Button("🧹 Clear Chat History", size="sm")
 
+    # Connect to your actual analyzer functions
     analyze_btn.click(
         analyzer.analyze_image,
         inputs=image_input,
         outputs=analysis_output
     )
-
+    
     chat_input.submit(
         lambda msg, chat: ("", chat + [(msg, analyzer.analyze_chat(msg))]),
         inputs=[chat_input, chatbot],
         outputs=[chat_input, chatbot]
     )
-
+    
     chat_btn.click(
         lambda msg, chat: ("", chat + [(msg, analyzer.analyze_chat(msg))]),
         inputs=[chat_input, chatbot],
         outputs=[chat_input, chatbot]
     )
+    
+    clear_btn.click(
+        lambda: None,
+        outputs=[chatbot]
+    )
 
-    clear_btn.click(lambda: None, outputs=[chatbot])
-
+    # Examples using your real analyzer
     gr.Examples(
         examples=[
-            ["examples/polluted_river1.jpg"],
-            ["examples/polluted_river2.jpg"]
+            ["https://drive.google.com/uc?export=view&id=1sCxcpacS5WkV5qVrhj8mcdq1JHyVyaEb"],
+            ["https://drive.google.com/uc?export=view&id=1WGcXwFhpbD1LrtbQ8E5IZZN3nEGfcwuN"]
         ],
         inputs=image_input,
         outputs=analysis_output,
         fn=analyzer.analyze_image,
-        cache_examples=torch.cuda.is_available(),  # Cache only if GPU available
-        label="Example Images"
+        cache_examples=True,
+        label="Try example images:"
     )
 
-demo.queue(max_size=2).launch()
+demo.launch()