Upload app.py

app.py

@@ -0,0 +1,144 @@
+# app.py
+import torch
+from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
+import gradio as gr
+from PIL import Image
+import re
+from typing import List, Tuple
+import os
+
+# Set cache directory (important for Hugging Face Spaces)
+CACHE_DIR = "model_cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+
+class RiverPollutionAnalyzer:
+    def __init__(self):
+        try:
+            # Load with caching and optimized device placement
+            self.processor = InstructBlipProcessor.from_pretrained(
+                "Salesforce/instructblip-vicuna-7b", cache_dir=CACHE_DIR
+            )
+
+            self.model = InstructBlipForConditionalGeneration.from_pretrained(
+                "Salesforce/instructblip-vicuna-7b",
+                cache_dir=CACHE_DIR,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                offload_folder="offload",
+                low_cpu_mem_usage=True,
+            )
+
+        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):
+        """Analyze river pollution with robust parsing"""
+        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
+
+Respond EXACTLY in this format:
+Pollutants: [comma separated list]
+Severity: [number]"""
+
+        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=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)
+
+    # ... [keep all your other methods unchanged] ...
+
+
+# Initialize analyzer with error handling
+try:
+    analyzer = RiverPollutionAnalyzer()
+    model_status = "Model loaded successfully!"
+except Exception as e:
+    analyzer = None
+    model_status = f"Model failed to load: {str(e)}"
+    print(model_status)
+
+
+# Create wrapper function for Gradio
+def analyze_image_wrapper(image):
+    if analyzer is None:
+        return (
+            f"⚠️ Error: {model_status}\nPlease try again later or use a smaller image."
+        )
+    return analyzer.analyze_image(image)
+
+
+# Gradio interface
+css = """
+/* [keep your existing CSS] */
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+    # [keep all your existing UI code]
+
+    analyze_btn.click(
+        analyze_image_wrapper, inputs=image_input, outputs=analysis_output
+    )
+
+    # Update examples to use wrapper
+    gr.Examples(
+        examples=[
+            [
+                "https://huggingface.co/spaces/Atharwaaah/SLCR-FLOWCODE-tarak.AI/resolve/main/polluted_river1.jpg"
+            ],
+            [
+                "https://huggingface.co/spaces/Atharwaaah/SLCR-FLOWCODE-tarak.AI/resolve/main/polluted_river2.jpg"
+            ],
+        ],
+        inputs=image_input,
+        outputs=analysis_output,
+        fn=analyze_image_wrapper,
+        cache_examples=True,
+        label="Try example images:",
+    )
+
+demo.launch()