Update app.py

app.py

@@ -5,6 +5,7 @@ from PIL import Image
 from transformers import AutoProcessor, AutoModelForVision2Seq
 import cv2
 from huggingface_hub import hf_hub_download
+import os
 from segment_anything import SamPredictor, sam_model_registry
 
 # India-specific constants
@@ -18,60 +19,83 @@ ANNUAL_MAINTENANCE_COST_PERCENT = 0.01  # 1% of installation cost
 LIFESPAN_YEARS = 25  # typical solar panel lifespan
 
 def load_models():
-    print("Loading models...")
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    # Load LLaVA model (float16 on GPU if available)
-    processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
-    model = AutoModelForVision2Seq.from_pretrained(
-        "llava-hf/llava-1.5-7b-hf", 
-        torch_dtype=torch.float16 if device.type=="cuda" else torch.float32
-    )
-    model.to(device)
-
-    # Download SAM checkpoint from your HF repo at runtime
-    repo_id = "kunkaran/sam_vit_h_4b8939.pth"  # Replace with your actual repo if different
-    filename = "sam_vit_h_4b8939.pth"          # Filename in that repo
-    print(f"Downloading checkpoint {filename} from repo {repo_id} ...")
-    sam_checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
-    print(f"Checkpoint downloaded to {sam_checkpoint_path}")
-
-    # Load SAM model
-    sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint_path)
-    sam.to(device)
-    sam_predictor = SamPredictor(sam)
-
-    print("Models loaded.")
-    return processor, model, sam_predictor, device
+    try:
+        print("Loading models...")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {device}")
+
+        # Load LLaVA model (float16 on GPU if available)
+        processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
+        model = AutoModelForVision2Seq.from_pretrained(
+            "llava-hf/llava-1.5-7b-hf", 
+            torch_dtype=torch.float16 if device.type=="cuda" else torch.float32
+        )
+        model.to(device)
+        print("LLaVA model loaded successfully")
+
+        # Download SAM checkpoint from your HF repo at runtime
+        repo_id = "kunkaran/sam_vit_h_4b8939.pth"  # Replace with your actual repo if different
+        filename = "sam_vit_h_4b8939.pth"          # Filename in that repo
+        
+        # Check if file already exists locally to avoid redundant downloads
+        local_path = os.path.join(os.getcwd(), filename)
+        if os.path.exists(local_path):
+            sam_checkpoint_path = local_path
+            print(f"Using existing checkpoint at {sam_checkpoint_path}")
+        else:
+            print(f"Downloading checkpoint {filename} from repo {repo_id} ...")
+            sam_checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
+            print(f"Checkpoint downloaded to {sam_checkpoint_path}")
+
+        # Load SAM model
+        sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint_path)
+        sam.to(device)
+        sam_predictor = SamPredictor(sam)
+        print("SAM model loaded successfully")
+
+        return processor, model, sam_predictor, device
+        
+    except Exception as e:
+        print(f"Error loading models: {str(e)}")
+        raise
 
 def segment_rooftop(image, sam_predictor):
-    if isinstance(image, Image.Image):
-        image_array = np.array(image)
-    else:
-        image_array = image
-
-    sam_predictor.set_image(image_array)
-    h, w = image_array.shape[:2]
-    input_point = np.array([[w//2, h//2]])
-    input_label = np.array([1])
-
-    masks, _, _ = sam_predictor.predict(
-        point_coords=input_point,
-        point_labels=input_label,
-        multimask_output=True
-    )
-    best_mask = masks[0]  # Choose first mask
-    return best_mask
+    try:
+        if isinstance(image, Image.Image):
+            image_array = np.array(image)
+        else:
+            image_array = image
+
+        # Ensure image is RGB if it's grayscale
+        if len(image_array.shape) == 2 or (len(image_array.shape) == 3 and image_array.shape[2] == 1):
+            image_array = cv2.cvtColor(image_array, cv2.COLOR_GRAY2RGB)
+
+        sam_predictor.set_image(image_array)
+        h, w = image_array.shape[:2]
+        input_point = np.array([[w//2, h//2]])
+        input_label = np.array([1])
+
+        masks, _, _ = sam_predictor.predict(
+            point_coords=input_point,
+            point_labels=input_label,
+            multimask_output=True
+        )
+        best_mask = masks[0]  # Choose first mask
+        return best_mask
+    except Exception as e:
+        print(f"Error in segmentation: {str(e)}")
+        # Return a blank mask as fallback
+        return np.zeros((image_array.shape[0], image_array.shape[1]), dtype=bool)
 
 def calculate_area(mask, image_width_meters=10.0):
     pixel_area = np.sum(mask)
     total_pixels = mask.shape[0] * mask.shape[1]
     image_area_m2 = image_width_meters * image_width_meters  # Approximate
     rooftop_area_m2 = (pixel_area / total_pixels) * image_area_m2
-    return rooftop_area_m2
+    return max(0.1, rooftop_area_m2)  # Ensure we don't return zero area
 
 def calculate_roi(area_m2):
-    num_panels = int(area_m2 / PANEL_SIZE * 0.8)
+    num_panels = max(1, int(area_m2 / PANEL_SIZE * 0.8))
     total_capacity_kw = num_panels * PANEL_CAPACITY
     daily_production_kwh = total_capacity_kw * SOLAR_IRRADIATION_INDIA * SOLAR_PANEL_EFFICIENCY
     annual_production_kwh = daily_production_kwh * 365
@@ -94,32 +118,43 @@ def calculate_roi(area_m2):
     }
 
 def analyze_rooftop(image, processor, model, sam_predictor, device):
-    mask = segment_rooftop(image, sam_predictor)
-    area_m2 = calculate_area(mask)
-    roi_data = calculate_roi(area_m2)
-
-    prompt = (
-        "Analyze this rooftop image for solar panel installation. "
-        "Comment on orientation, obstructions, and ideal panel placement."
-    )
-
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
-    output = model.generate(**inputs, max_new_tokens=300)
-    llava_analysis = processor.decode(output[0], skip_special_tokens=True)
-
-    image_array = np.array(image)
-    visualization = image_array.copy()
-    mask_overlay = np.zeros_like(visualization)
-    if visualization.ndim == 2:  # grayscale image
-        visualization = cv2.cvtColor(visualization, cv2.COLOR_GRAY2RGB)
-    mask_overlay[:,:,0] = (mask * 255).astype(np.uint8)
-    visualization = cv2.addWeighted(visualization.astype(np.uint8), 0.7, mask_overlay, 0.3, 0)
-    visualization = Image.fromarray(visualization)
-
-    results_str = f"""
+    try:
+        mask = segment_rooftop(image, sam_predictor)
+        area_m2 = calculate_area(mask)
+        roi_data = calculate_roi(area_m2)
+
+        prompt = (
+            "As a solar engineering expert in India, analyze this rooftop for solar installation potential. "
+            "Consider: 1) Roof orientation and tilt for optimal sun exposure in Indian latitudes, "
+            "2) Potential obstructions like chimneys, water tanks, or shadows from nearby structures, "
+            "3) Structural integrity and recommended panel layout for monsoon resistance, "
+            "4) Regional factors like dust accumulation and temperature impact on efficiency, "
+            "5) Space utilization for maximizing energy generation with Indian solar irradiation patterns. "
+            "Provide a detailed assessment for this specific rooftop."
+        )
+
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
+        output = model.generate(**inputs, max_new_tokens=300)
+        llava_analysis = processor.decode(output[0], skip_special_tokens=True)
+
+        image_array = np.array(image)
+        visualization = image_array.copy()
+        
+        # Ensure visualization array is RGB
+        if visualization.ndim == 2:  # grayscale image
+            visualization = cv2.cvtColor(visualization, cv2.COLOR_GRAY2RGB)
+        elif visualization.shape[2] == 4:  # RGBA image
+            visualization = visualization[:,:,:3]
+            
+        mask_overlay = np.zeros_like(visualization)
+        mask_overlay[:,:,0] = (mask * 255).astype(np.uint8)
+        visualization = cv2.addWeighted(visualization.astype(np.uint8), 0.7, mask_overlay, 0.3, 0)
+        visualization = Image.fromarray(visualization)
+
+        results_str = f"""
 ## Rooftop Analysis Results
 ### Technical Assessment
 - Usable Rooftop Area: {roi_data['rooftop_area_m2']} m²
@@ -135,8 +170,13 @@ def analyze_rooftop(image, processor, model, sam_predictor, device):
 {llava_analysis}
 """
 
-    return visualization, results_str
+        return visualization, results_str
+    except Exception as e:
+        error_message = f"Error during analysis: {str(e)}"
+        print(error_message)
+        return Image.fromarray(np.zeros((400, 400, 3), dtype=np.uint8)), f"## Error\n{error_message}"
 
+print("Initializing application...")
 # Load models once (synchronously)
 processor, model, sam_predictor, device = load_models()
 
@@ -150,7 +190,9 @@ demo = gr.Interface(
     ],
     title="🔆 Solar Rooftop Analyzer - India Edition",
     description="Upload a satellite image of a rooftop to get detailed solar potential analysis with India-specific ROI calculations.",
+    examples=["sample1.jpg", "sample2.jpg"] if os.path.exists("sample1.jpg") else None
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    print("Starting Gradio server...")
+    demo.launch()