import gradio as gr
import numpy as np
from preprocess import preprocess_image_from_array


def process_image(image):
    """
    Process an uploaded image and return the preprocessed version.
    
    Args:
        image: PIL Image or numpy array from Gradio
        
    Returns:
        Preprocessed image as numpy array
    """
    if image is None:
        return None
    
    try:
        # Convert PIL Image to numpy array if needed
        if hasattr(image, 'mode'):
            # PIL Image
            img_array = np.array(image)
        else:
            # Already a numpy array
            img_array = image
        
        # Ensure RGB format
        if len(img_array.shape) == 2:
            # Grayscale, convert to RGB
            img_array = np.stack([img_array] * 3, axis=-1)
        elif img_array.shape[2] == 4:
            # RGBA, convert to RGB
            img_array = img_array[:, :, :3]
        
        # Process the image
        processed = preprocess_image_from_array(img_array)
        
        return processed
    except Exception as e:
        print(f"Erreur lors du traitement: {e}")
        import traceback
        traceback.print_exc()
        return None


# Create Gradio interface
with gr.Blocks(title="Preprocessing d'Échiquier") as demo:
    gr.Markdown("""
    # Preprocessing d'Images d'Échiquier
    
    Cette application applique le preprocessing d'images d'échiquier pour convertir 
    une photo prise sous un angle arbitraire en une projection 2D.
    
    Uploadez une image d'échiquier et obtenez la version préprocessée.
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Image d'entrée", type="numpy")
            process_btn = gr.Button("Traiter l'image", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Image préprocessée", type="numpy")
    
    process_btn.click(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_image]
    )
    
    # Auto-process when image is uploaded
    input_image.change(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_image]
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)