import gradio as gr
import numpy as np
from PIL import Image

from segmentation import segment_image, crop_dog, visualize_segmentation
from classification import classify

# config
pre_scale_size = (2048, 2048)


def classify_image_with_cropping(original_image, pre_segment):
    if isinstance(original_image, np.ndarray):
        original_image = Image.fromarray(original_image)  # Convert ndarray to PIL Image

    # 1. Pre-scale
    if original_image.width > pre_scale_size[0] or original_image.height > pre_scale_size[1]:
        original_image.thumbnail(pre_scale_size, Image.LANCZOS)

    # 1. Segment the image
    if pre_segment:
        print("Segmenting...")
        segmented_image, mask = segment_image(original_image)

        if mask is not None:
            # 2. Crop to the dog (if found)
            print("Cropping...")
            visualised_image = visualize_segmentation(original_image, mask)
            cropped_image = crop_dog(segmented_image, mask)
        else:
            print(f"Failed segmentation, using original image")
            visualised_image = None
            cropped_image = original_image
    else:
        visualised_image = None
        cropped_image = original_image

    # 3. Preprocess and classify the cropped image
    print("Running classifier...")
    predicted_class_idx, predicted_label = classify(cropped_image)

    print("Done.")
    return visualised_image, cropped_image, predicted_label


iface = gr.Interface(
    fn=classify_image_with_cropping,
    inputs=[gr.Image(type="pil"),
            gr.Checkbox(label="Try to isolate dog (pre-segmentation)", value=True)],
    outputs=[gr.Image(type="pil", label="Segmented image"),
             gr.Image(type="pil", label="Predicted image"),
             gr.Textbox(label="Predicated class")]
)

iface.launch()