import gradio as gr
import cv2
import numpy as np
from utils.preprocessing import ImageProcessor

# Initialize processor
processor = ImageProcessor("models/best.pt")

def process_image(input_image):
    if input_image is None:
        raise gr.Error("Please upload an image first!")
    
    # Convert Gradio Image to bytes
    _, img_bytes = cv2.imencode(".png", input_image)
    
    # Process image
    results = processor.process_image(img_bytes.tobytes())
    
    # Format outputs
    return {
        class_name: (mask * 255).astype(np.uint8)
        for class_name, mask in results.items()
    }

# Gradio interface
with gr.Blocks(title="Fashion Segmenter") as demo:
    gr.Markdown("# 🧥 Fashion Item Segmenter")
    
    with gr.Row():
        input_image = gr.Image(label="Upload Clothing Image", type="numpy")
        output_gallery = gr.Gallery(label="Segmented Items", columns=2)
    
    with gr.Row():
        run_btn = gr.Button("Process Image", variant="primary")
        examples = gr.Examples(
            examples=["sample1.jpg", "sample2.jpg"],
            inputs=[input_image],
            label="Example Images"
        )

    run_btn.click(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_gallery],
        show_progress=True
    )

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