app.py

import gradio as gr
import torch
import numpy as np
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
SOLAR_IRRADIATION_INDIA = 5.0  # kWh/m²/day (average for India)
SOLAR_PANEL_EFFICIENCY = 0.18  # 18% efficiency for modern panels
PANEL_SIZE = 1.7  # m² (typical solar panel size)
PANEL_CAPACITY = 0.340  # kW per panel (typical 340W panel)
INSTALLATION_COST_PER_KW = 45000  # INR per kW
ELECTRICITY_RATE = 8  # INR per kWh
ANNUAL_MAINTENANCE_COST_PERCENT = 0.01  # 1% of installation cost
LIFESPAN_YEARS = 25  # typical solar panel lifespan

def load_models():
    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):
    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 max(0.1, rooftop_area_m2)  # Ensure we don't return zero area

def calculate_roi(area_m2):
    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
    total_installation_cost = total_capacity_kw * INSTALLATION_COST_PER_KW
    annual_savings = annual_production_kwh * ELECTRICITY_RATE
    annual_maintenance = total_installation_cost * ANNUAL_MAINTENANCE_COST_PERCENT
    net_annual_savings = annual_savings - annual_maintenance
    simple_payback_years = total_installation_cost / net_annual_savings if net_annual_savings > 0 else float('inf')
    lifetime_savings = net_annual_savings * LIFESPAN_YEARS - total_installation_cost

    return {
        "rooftop_area_m2": round(area_m2, 2),
        "num_panels": num_panels,
        "total_capacity_kw": round(total_capacity_kw, 2),
        "annual_production_kwh": round(annual_production_kwh, 2),
        "installation_cost_inr": round(total_installation_cost, 2),
        "annual_savings_inr": round(annual_savings, 2),
        "payback_years": round(simple_payback_years, 2) if simple_payback_years != float('inf') else "N/A",
        "lifetime_savings_inr": round(lifetime_savings, 2),
    }

def analyze_rooftop(image, processor, model, sam_predictor, device):
    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²
- Potential Solar Panel Count: {roi_data['num_panels']} panels
- Total System Capacity: {roi_data['total_capacity_kw']} kW
- Annual Energy Production: {roi_data['annual_production_kwh']} kWh
### Financial Analysis (INR)
- Installation Cost: ₹{roi_data['installation_cost_inr']:,.2f}
- Annual Savings: ₹{roi_data['annual_savings_inr']:,.2f}
- Payback Period: {roi_data['payback_years']} years
- 25-Year Lifetime Savings: ₹{roi_data['lifetime_savings_inr']:,.2f}
### Expert Analysis
{llava_analysis}
"""

        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()

# Gradio interface
demo = gr.Interface(
    fn=lambda img: analyze_rooftop(img, processor, model, sam_predictor, device),
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(label="Segmentation Visualization"),
        gr.Markdown(label="Analysis Results")
    ],
    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__":
    print("Starting Gradio server...")
    demo.launch()