Update app.py

app.py

@@ -1,24 +1,179 @@
+!pip install -q transformers accelerate bitsandbytes gradio torch pillow
+
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    BlipProcessor,
+    BlipForConditionalGeneration,
+    BitsAndBytesConfig
+)
 import gradio as gr
 from PIL import Image
-from transformers import BlipProcessor, BlipForConditionalGeneration
-import torch
+import re
+from typing import List, Tuple
 
-# Load a smaller, GPU-compatible model
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-
-def generate_caption(image):
-    inputs = processor(images=image, return_tensors="pt")
-    out = model.generate(**inputs)
-    caption = processor.decode(out[0], skip_special_tokens=True)
-    return caption
-
-demo = gr.Interface(
-    fn=generate_caption,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Image Caption Generator",
-    description="Upload an image and get a caption using BLIP base model."
+# Configuration for 4-bit quantization
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True
 )
 
-demo.launch()
+class RiverPollutionAnalyzer:
+    def __init__(self):
+        try:
+            # Initialize BLIP for image captioning
+            self.blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+            self.blip_model = BlipForConditionalGeneration.from_pretrained(
+                "Salesforce/blip-image-captioning-base",
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            
+            # Initialize FLAN-T5-XL for text analysis
+            self.tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-xl")
+            self.model = AutoModelForSeq2SeqLM.from_pretrained(
+                "google/flan-t5-xl",
+                device_map="auto",
+                quantization_config=quant_config
+            )
+            
+        except Exception as e:
+            raise RuntimeError(f"Model loading failed: {str(e)}")
+
+        self.pollutants = [
+            "plastic waste", "chemical foam", "industrial discharge",
+            "sewage water", "oil spill", "organic debris",
+            "construction waste", "medical waste", "floating trash",
+            "algal bloom", "toxic sludge", "agricultural runoff"
+        ]
+
+        self.severity_descriptions = {
+            1: "Minimal pollution - Slightly noticeable",
+            2: "Minor pollution - Small amounts visible",
+            3: "Moderate pollution - Clearly visible",
+            4: "Significant pollution - Affecting water quality",
+            5: "Heavy pollution - Obvious environmental impact",
+            6: "Severe pollution - Large accumulation",
+            7: "Very severe pollution - Major ecosystem impact",
+            8: "Extreme pollution - Dangerous levels",
+            9: "Critical pollution - Immediate action needed",
+            10: "Disaster level - Ecological catastrophe"
+        }
+
+    def analyze_image(self, image):
+        """Two-step analysis: BLIP captioning + FLAN-T5 analysis"""
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+        
+        try:
+            # Step 1: Generate image caption with BLIP
+            inputs = self.blip_processor(image, return_tensors="pt").to(self.blip_model.device, torch.float16)
+            caption = self.blip_model.generate(**inputs, max_new_tokens=100)[0]
+            caption = self.blip_processor.decode(caption, skip_special_tokens=True)
+            
+            # Step 2: Analyze caption with FLAN-T5
+            prompt = f"""Analyze this river scene: '{caption}'
+1. List visible pollutants from: {self.pollutants}
+2. Estimate severity (1-10)
+
+Respond EXACTLY as:
+Pollutants: [comma separated list]
+Severity: [number]"""
+            
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+            outputs = self.model.generate(**inputs, max_new_tokens=200)
+            analysis = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            pollutants, severity = self._parse_response(analysis)
+            return self._format_analysis(pollutants, severity)
+            
+        except Exception as e:
+            return f"⚠️ Analysis failed: {str(e)}"
+
+    # [Keep all your existing parsing/formatting methods unchanged]
+    def _parse_response(self, analysis: str) -> Tuple[List[str], int]:
+        """Same parsing logic as before"""
+        # ... (unchanged from your original code) ...
+    
+    def _calculate_severity(self, pollutants: List[str]) -> int:
+        """Same severity calculation"""
+        # ... (unchanged from your original code) ...
+    
+    def _format_analysis(self, pollutants: List[str], severity: int) -> str:
+        """Same formatting"""
+        # ... (unchanged from your original code) ...
+
+    def analyze_chat(self, message: str) -> str:
+        """Same chat handler"""
+        # ... (unchanged from your original code) ...
+
+# Initialize with error handling
+try:
+    analyzer = RiverPollutionAnalyzer()
+    model_status = "✅ Models loaded successfully"
+except Exception as e:
+    analyzer = None
+    model_status = f"❌ Model loading failed: {str(e)}"
+
+# Gradio Interface (unchanged layout from your original)
+css = """
+/* [Keep your existing CSS] */
+"""
+
+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}")
+
+    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")
+                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)
+                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")
+
+    # Connect functions
+    analyze_btn.click(
+        analyzer.analyze_image if analyzer else lambda x: "Model not loaded",
+        inputs=image_input,
+        outputs=analysis_output
+    )
+
+    # [Keep all other UI event handlers unchanged]
+
+    # Update examples to use local files
+    gr.Examples(
+        examples=[
+            ["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:"
+    )
+
+# Launch with queue for stability
+demo.queue(max_size=3).launch()