import gradio as gr
import numpy as np
import cv2
from PIL import Image
import tempfile
from infer import infer_image
 
def create_output_image(result):
    """Convert model output to proper image format"""
    if len(result.shape) == 2:  # Grayscale
        result = cv2.cvtColor(result, cv2.COLOR_GRAY2RGB)
    elif result.shape[2] == 1:  # Single channel
        result = cv2.cvtColor(result, cv2.COLOR_GRAY2RGB)
    elif result.shape[2] == 3:  # BGR
        result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
    return result
 
def process_image(input_image):
    try:
        # Convert PIL to numpy
        image_np = np.array(input_image)
        
        # Process through model
        result = infer_image(image_np)
        
        # Handle error image case
        if result.mean() == 0 and "error" in str(result).lower():
            raise RuntimeError("Model processing failed")
        
        # Create proper output
        output_image = create_output_image(result)
        
        # Save to temp file
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
            filepath = tmp_file.name
            Image.fromarray(output_image).save(filepath, quality=95)
            return filepath
            
    except Exception as e:
        # Create error image
        error_img = np.zeros((512, 512, 3), dtype=np.uint8)
        cv2.putText(error_img, "Processing Error", (50, 100), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)
        cv2.putText(error_img, str(e)[:100], (50, 150), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
        
        # Save error image
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
            filepath = tmp_file.name
            cv2.imwrite(filepath, cv2.cvtColor(error_img, cv2.COLOR_RGB2BGR))
            return filepath
 
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.File(label="Processed Floorplan"),
    title="Floorplan Processor",
    description="Upload an image to generate a processed floorplan",
    allow_flagging="never"
)
 
iface.launch(server_name="0.0.0.0", server_port=7860)