import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import numpy as np
from landingai.common import decode_bitmap_rle
import cv2
import os

ENDPOINT_ID = os.environ.get("crosswalk_detection_model_endpoint_ID_Nitish_Borthakur")
API_KEY = os.environ.get("crosswalk_detection_model_API_key_Nitish_Borthakur")
API_URL = f"https://predict.app.landing.ai/inference/v1/predict?endpoint_id={ENDPOINT_ID}"

def predict_from_landinglens(image_path):
    # Load and keep original imag
    original_img = Image.open(image_path).convert("RGB")
    img_array = np.array(original_img)

    # Get image dimensions
    height, width = img_array.shape[:2]
    total_pixels = height * width

    # Prepare for API
    buffered = BytesIO()
    original_img.save(buffered, format="JPEG")
    img_bytes = buffered.getvalue()

    files = {"file": (image_path, img_bytes, "image/jpeg")}
    headers = {"apikey": API_KEY}

    try:
        response = requests.post(API_URL, files=files, headers=headers)
        if response.status_code == 503:
            return "Service temporarily unavailable. Please try again later."
        response.raise_for_status()
        prediction = response.json()

        if "predictions" not in prediction or not prediction.get("predictions"):
            print("No 'predictions' key found or it's empty.")
            return "Error: No 'predictions' found."

        bitmaps = prediction["predictions"]["bitmaps"]
        masked_images = []
        coverage_info = []

        for i, (bitmap_id, bitmap_data) in enumerate(bitmaps.items()):
            try:
                # Decode mask
                mask = decode_bitmap_rle(bitmap_data["bitmap"])
                if isinstance(mask, list):
                    mask = np.array(mask)

                # Reshape mask to match image dimensions
                mask = mask.reshape(prediction["predictions"]["imageHeight"],
                                  prediction["predictions"]["imageWidth"])

                # Calculate area coverage
                mask_area = np.sum(mask > 0)
                coverage_percentage = (mask_area / total_pixels) * 100
                label_name = bitmap_data.get("label_name", f"Mask {i}")
                coverage_info.append(f"{label_name}: {coverage_percentage:.2f}%")

                # Create colored overlay
                colored_mask = np.zeros_like(img_array)
                colored_mask[mask > 0] = [255, 0, 0]  # Red overlay for mask

                # Combine original image with colored mask
                alpha = 0.5  # Transparency of the overlay
                combined = cv2.addWeighted(img_array, 1, colored_mask, alpha, 0)

                # Convert to PIL Image
                masked_image = Image.fromarray(combined)
                masked_images.append(masked_image)

            except Exception as e:
                print(f"Error processing mask {i}: {e}")
                continue

        return masked_images, "\n".join(coverage_info)

    except requests.exceptions.RequestException as e:
        print(f"API Error: {e}")
        return f"API Error: {e}"

iface = gr.Interface(
    fn=predict_from_landinglens,
    inputs=gr.Image(type="filepath"),
    outputs=[
        gr.Gallery(format="png"),
        gr.Textbox(label="Area of each mask in the image")
    ],
    title="Crosswalk detection app",
)

if __name__ == "__main__":
    iface.launch()