Update app.py

app.py

@@ -1,5 +1,3 @@
-
-
 import torch
 from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
 import gradio as gr
@@ -7,16 +5,19 @@ from PIL import Image
 import re
 from typing import List, Tuple
 
+# Configuration
+MODEL_NAME = "Salesforce/instructblip-flan-t5-xl"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TORCH_DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
+
 class RiverPollutionAnalyzer:
     def __init__(self):
-        # Initialize model with 4-bit quantization
-        self.processor = InstructBlipProcessor.from_pretrained("Salesforce/instructblip-vicuna-7b")
+        # Initialize processor and model
+        self.processor = InstructBlipProcessor.from_pretrained(MODEL_NAME)
         self.model = InstructBlipForConditionalGeneration.from_pretrained(
-            "Salesforce/instructblip-vicuna-7b",
-            device_map="auto",
-            torch_dtype=torch.float16,
-            load_in_4bit=True
-        )
+            MODEL_NAME,
+            torch_dtype=TORCH_DTYPE
+        ).to(DEVICE)
 
         self.pollutants = [
             "plastic waste", "chemical foam", "industrial discharge",
@@ -55,14 +56,14 @@ Severity: [number]"""
             images=image,
             text=prompt,
             return_tensors="pt"
-        ).to("cuda", torch.float16)
+        ).to(DEVICE, TORCH_DTYPE)
 
         with torch.no_grad():
             outputs = self.model.generate(
                 **inputs,
                 max_new_tokens=200,
-                temperature=0.5,
-                top_p=0.85,
+                temperature=0.7,
+                top_p=0.9,
                 do_sample=True
             )
 
@@ -70,6 +71,15 @@ Severity: [number]"""
         pollutants, severity = self._parse_response(analysis)
         return self._format_analysis(pollutants, severity)
 
+    def analyze_chat(self, message):
+        """Handle chat questions about pollution"""
+        if "severity" in message.lower():
+            return "Severity levels range from 1 (minimal) to 10 (disaster). The analyzer automatically detects the appropriate level."
+        elif "pollutant" in message.lower():
+            return f"Detectable pollutants: {', '.join(self.pollutants)}"
+        else:
+            return "I can answer questions about pollution severity levels and detectable pollutants."
+
     def _parse_response(self, analysis: str) -> Tuple[List[str], int]:
         """Robust parsing of model response"""
         pollutants = []
@@ -77,29 +87,24 @@ Severity: [number]"""
 
         # Extract pollutants
         pollutant_match = re.search(
-            r'(?i)(pollutants?|contaminants?)[:\s]*\[?(.*?)(?:\]|Severity|severity|$)',
-            analysis
+            r'Pollutants:\s*\[?(.*?)\]?',
+            analysis, re.IGNORECASE
         )
-
         if pollutant_match:
-            pollutants_str = pollutant_match.group(2).strip()
+            pollutants_str = pollutant_match.group(1).strip()
             pollutants = [
-                p.strip().lower()
-                for p in re.split(r'[,;]|\band\b', pollutants_str)
+                p.strip().lower() 
+                for p in re.split(r'[,;]', pollutants_str)
                 if p.strip().lower() in self.pollutants
             ]
 
         # Extract severity
         severity_match = re.search(
-            r'(?i)(severity|level)[:\s]*(\d{1,2})',
-            analysis
+            r'Severity:\s*(\d{1,2})',
+            analysis, re.IGNORECASE
         )
-
         if severity_match:
-            try:
-                severity = min(max(int(severity_match.group(2)), 1), 10)
-            except:
-                severity = self._calculate_severity(pollutants)
+            severity = min(max(int(severity_match.group(1)), 1), 10)
         else:
             severity = self._calculate_severity(pollutants)
 
@@ -116,7 +121,6 @@ Severity: [number]"""
             "plastic waste": 1.5, "construction waste": 1.5, "algal bloom": 1.5,
             "agricultural runoff": 1.5, "floating trash": 1, "organic debris": 1
         }
-
         avg_weight = sum(weights.get(p, 1) for p in pollutants) / len(pollutants)
         return min(10, max(1, round(avg_weight * 3)))
 
@@ -127,7 +131,7 @@ Severity: [number]"""
 {self.severity_descriptions.get(severity, '')}"""
 
         pollutants_list = "\n🔍 No pollutants detected" if not pollutants else "\n".join(
-            f"{i}. {p.capitalize()}" for i, p in enumerate(pollutants[:5], 1))
+            f"• {p.capitalize()}" for p in pollutants[:8])
 
         return f"""🌊 River Pollution Analysis 🌊
 {pollutants_list}
@@ -136,75 +140,66 @@ Severity: [number]"""
 # Initialize analyzer
 analyzer = RiverPollutionAnalyzer()
 
-import gradio as gr
-
- # Import your actual analyzer
-
+# Gradio Interface
 css = """
-/* (Keep all your CSS styles) */
+.header { text-align: center; margin-bottom: 20px; }
+.header h1 { font-size: 2.2rem; margin-bottom: 0; }
+.header h3 { font-size: 1.1rem; font-weight: normal; margin-top: 0.5rem; }
+.side-by-side { display: flex; gap: 20px; }
+.left-panel, .right-panel { flex: 1; }
+.analysis-box { border: 1px solid #e0e0e0; border-radius: 8px; padding: 15px; margin-top: 20px; }
+.chat-container { border: 1px solid #e0e0e0; border-radius: 8px; padding: 15px; height: 100%; }
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     with gr.Column(elem_classes="header"):
         gr.Markdown("# 🌍 River Pollution Analyzer")
-        gr.Markdown("### AI-powered water pollution detection")
+        gr.Markdown("### AI-powered water quality assessment")
 
     with gr.Row(elem_classes="side-by-side"):
-        # Left Panel
+        # Image Analysis Panel
         with gr.Column(elem_classes="left-panel"):
+            gr.Markdown("### 📸 Image Analysis")
             with gr.Group():
                 image_input = gr.Image(type="pil", label="Upload River Image", height=300)
-                analyze_btn = gr.Button("🔍 Analyze Pollution", variant="primary")
-            
+                analyze_btn = gr.Button("🔍 Analyze", variant="primary")
             with gr.Group(elem_classes="analysis-box"):
-                gr.Markdown("### 📊 Analysis report")
                 analysis_output = gr.Markdown()
 
-        # Right Panel
+        # Chat Panel
         with gr.Column(elem_classes="right-panel"):
+            gr.Markdown("### 💬 Pollution Q&A")
             with gr.Group(elem_classes="chat-container"):
-                chatbot = gr.Chatbot(label="Pollution Analysis Q&A", height=400)
+                chatbot = gr.Chatbot(height=350)
                 with gr.Row():
-                    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 History", size="sm")
+                    chat_input = gr.Textbox(placeholder="Ask about pollution...", show_label=False)
+                    chat_btn = gr.Button("Send", variant="secondary")
+                clear_btn = gr.Button("Clear Chat")
 
-    # Connect to your actual analyzer functions
+    # Event handlers
     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]
-    )
 
-    # Examples using your real analyzer
+    def respond(message, chat_history):
+        response = analyzer.analyze_chat(message)
+        chat_history.append((message, response))
+        return "", chat_history
+
+    chat_input.submit(respond, [chat_input, chatbot], [chat_input, chatbot])
+    chat_btn.click(respond, [chat_input, chatbot], [chat_input, chatbot])
+    clear_btn.click(lambda: None, None, chatbot, queue=False)
+
+    # Examples
     gr.Examples(
-        examples=[
-            ["https://drive.google.com/uc?export=view&id=1sCxcpacS5WkV5qVrhj8mcdq1JHyVyaEb"],
-            ["https://drive.google.com/uc?export=view&id=1WGcXwFhpbD1LrtbQ8E5IZZN3nEGfcwuN"]
-        ],
+        examples=[["examples/pollution1.jpg"], ["examples/pollution2.jpg"]],
         inputs=image_input,
         outputs=analysis_output,
         fn=analyzer.analyze_image,
         cache_examples=True,
-        label="Try example images:"
+        label="Example Images"
     )
 
 demo.launch()