import gradio as gr
import geoai
import torch

model_path = "best_model.pth"
output_mask = "prediction.tif"
output_vector = "prediction.geojson"

def detect_objects(input_image):
    input_image.save("input.tif")

    geoai.object_detection(
        image_path="input.tif",
        output_path=output_mask,
        model_path=model_path,
        window_size=512,
        overlap=256,
        confidence_threshold=0.5,
        batch_size=4,
        num_channels=3,
    )

    geoai.orthogonalize(output_mask, output_vector, epsilon=2)

    return output_vector

gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil", label="Upload GeoTIFF"),
    outputs="file",
    title="Object Detection with GeoAI Mask R-CNN",
    description="Upload a 512x512 GeoTIFF to detect objects using a custom-trained Mask R-CNN model."
).launch()