Update app.py

app.py

@@ -1,11 +1,16 @@
+# app.py
 import torch
 from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration, BitsAndBytesConfig
 import gradio as gr
 from PIL import Image
 import re
+import os
 from typing import List, Tuple
 
-# Configuration for 4-bit quantization
+# Create model cache directory
+os.makedirs("model_cache", exist_ok=True)
+
+# 4-bit quantization config
 quant_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16,
@@ -16,7 +21,7 @@ quant_config = BitsAndBytesConfig(
 class RiverPollutionAnalyzer:
     def __init__(self):
         try:
-            # Initialize InstructBLIP-FLAN-T5-XL with 4-bit quantization
+            # Load InstructBLIP-FLAN-T5-XL with 4-bit quantization
             self.processor = InstructBlipProcessor.from_pretrained(
                 "Salesforce/instructblip-flan-t5-xl",
                 cache_dir="model_cache"
@@ -28,8 +33,12 @@ class RiverPollutionAnalyzer:
                 torch_dtype=torch.float16,
                 cache_dir="model_cache"
             )
+            self.model_loaded = True
+            self.status = "✅ Model loaded successfully"
         except Exception as e:
-            raise RuntimeError(f"Model loading failed: {str(e)}")
+            self.model_loaded = False
+            self.status = f"❌ Model loading failed: {str(e)}"
+            print(self.status)
 
         self.pollutants = [
             "plastic waste", "chemical foam", "industrial discharge",
@@ -52,15 +61,17 @@ class RiverPollutionAnalyzer:
         }
 
     def analyze_image(self, image):
-        """Analyze river pollution with robust parsing"""
+        if not self.model_loaded:
+            return "Model not loaded. Please check logs."
+            
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image)
 
-        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]
-2. Estimate pollution severity from 1-10
+        prompt = """Analyze this river pollution and list:
+1. Visible pollutants 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]
+2. Severity estimate (1-10)
 
-Respond EXACTLY in this format:
+Respond EXACTLY as:
 Pollutants: [comma separated list]
 Severity: [number]"""
 
@@ -71,22 +82,19 @@ Severity: [number]"""
                 return_tensors="pt"
             ).to(self.model.device)
 
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=200,
-                    temperature=0.5,
-                    top_p=0.85,
-                    do_sample=True
-                )
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=150,  # Reduced for memory
+                temperature=0.5,
+                top_p=0.85
+            )
 
             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 failed: {str(e)}"
+            return f"⚠️ Analysis error: {str(e)}"
 
-    # [Keep your existing parsing/formatting methods]
     def _parse_response(self, analysis: str) -> Tuple[List[str], int]:
         pollutants = []
         severity = 3
@@ -96,21 +104,16 @@ Severity: [number]"""
             r'(?i)(pollutants?|contaminants?)[:\s]*\[?(.*?)(?:\]|Severity|severity|$)',
             analysis
         )
-
         if pollutant_match:
             pollutants_str = pollutant_match.group(2).strip()
             pollutants = [
-                p.strip().lower()
+                p.strip().lower() 
                 for p in re.split(r'[,;]|\band\b', pollutants_str)
                 if p.strip().lower() in self.pollutants
             ]
 
         # 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)
@@ -124,14 +127,12 @@ Severity: [number]"""
     def _calculate_severity(self, pollutants: List[str]) -> int:
         if not pollutants:
             return 1
-
         weights = {
             "medical waste": 3, "toxic sludge": 3, "oil spill": 2.5,
             "chemical foam": 2, "industrial discharge": 2, "sewage water": 2,
             "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)))
 
@@ -147,21 +148,8 @@ Severity: [number]"""
 {pollutants_list}
 {severity_bar}"""
 
-    def analyze_chat(self, message: str) -> str:
-        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 or upload an image for analysis."
-        elif "pollution" in message.lower():
-            return "Common river pollutants include: plastic waste, chemical foam, industrial discharge, sewage water, and oil spills."
-        else:
-            return "I can answer questions about river pollution. Try asking about pollution types or upload an image for analysis."
-
-# Initialize with error handling
-try:
-    analyzer = RiverPollutionAnalyzer()
-    model_status = "✅ Model loaded successfully"
-except Exception as e:
-    analyzer = None
-    model_status = f"❌ Model loading failed: {str(e)}"
+# Initialize analyzer
+analyzer = RiverPollutionAnalyzer()
 
 css = """
 .header {
@@ -191,66 +179,58 @@ css = """
     border-radius: 10px;
     height: 100%;
 }
+.dark .analysis-box, .dark .chat-container {
+    background: #2a2a2a;
+    border-color: #444;
+}
 """
 
 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"### {model_status}")
+        gr.Markdown(f"### {analyzer.status}")
 
     with gr.Row(elem_classes="side-by-side"):
+        # Left Panel
         with gr.Column(elem_classes="left-panel"):
             with gr.Group():
                 image_input = gr.Image(type="pil", label="Upload River Image", height=300)
                 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 Analysis Q&A", height=400)
+                gr.Markdown("### 💬 Pollution Q&A")
+                chatbot = gr.Chatbot(height=400)
                 with gr.Row():
                     chat_input = gr.Textbox(
-                        placeholder="Ask about pollution sources...",
-                        label="Your Question",
+                        placeholder="Ask about pollution types...",
                         container=False,
                         scale=5
                     )
-                    chat_btn = gr.Button("💬 Ask", variant="secondary", scale=1)
-                clear_btn = gr.Button("🧹 Clear Chat History", size="sm")
+                    chat_btn = gr.Button("Send", variant="secondary", scale=1)
+                clear_btn = gr.Button("Clear Chat", size="sm")
 
     analyze_btn.click(
-        analyzer.analyze_image if analyzer else lambda x: "Model not loaded",
+        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])
-
     gr.Examples(
         examples=[
-            ["https://huggingface.co/spaces/atharwaah1work/tarak.AI/resolve/main/polluted_river1.jpg"],
-            ["https://huggingface.co/spaces/atharwaah1work/tarak.AI/resolve/main/polluted_river2.jpg"]
+            ["examples/polluted_river1.jpg"],
+            ["examples/polluted_river2.jpg"]
         ],
         inputs=image_input,
         outputs=analysis_output,
-        fn=analyzer.analyze_image if analyzer else lambda x: "Model not loaded",
-        cache_examples=True,
-        label="Try example images:"
+        fn=analyzer.analyze_image,
+        cache_examples=False,  # Disabled for free tier
+        label="Example Images"
     )
 
-demo.queue(max_size=3).launch()
+demo.queue(max_size=2).launch()