import numpy as np
import cv2
import torch
import albumentations as albu
from pylab import imshow
import matplotlib.pyplot as plt
from diffusers import StableDiffusionInpaintPipeline
from PIL import Image
from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad
from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image
from cloths_segmentation.pre_trained_models import create_model
import gradio as gr

# Load Cloth Segmentation Model (Ensure this is available)
try:
    model = create_model("Unet_2020-10-30")
    model.eval()
except Exception as e:
    raise RuntimeError(f"Error loading segmentation model: {e}")

# Load Inpainting Model (Without CUDA)
try:
    pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
    
except Exception as e:
    raise RuntimeError(f"Error loading inpainting model: {e}")


def load_and_preprocess_image(image_path):
    image = load_rgb(image_path)
    # Pad the image to be divisible by 32
    padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)  
    transform = albu.Compose([albu.Normalize(p=1)], p=1)
    x = transform(image=padded_image)["image"]  # Use padded image here
    x = torch.unsqueeze(tensor_from_rgb_image(x), 0)
    return x, image, pads  # Return original image and padding info for later unpadding

def segment_cloth(image_tensor, model, pads):
    with torch.no_grad():
        prediction = model(image_tensor)[0][0]
    mask = (prediction > 0).cpu().numpy().astype(np.uint8)
    mask = unpad(mask, pads)
    return mask

def perform_inpainting(image_path, mask_path, prompt):
    image = Image.open(image_path)
    mask_image = Image.open(mask_path).convert("L") 
    mask_image = mask_image.resize(image.size)
    output_image = pipe(prompt=prompt, image=image, mask_image=mask_image).images[0]
    return output_image

def resize_and_upscale(image, new_width, new_height):
    resized_img = cv2.resize(np.array(image), (new_width, new_height), interpolation=cv2.INTER_CUBIC)
    return Image.fromarray(resized_img)

def image_segmentation_and_inpainting(image, prompt="Chinese Red and Golder Armor"):
    try:
        pil_image = Image.fromarray(image.astype('uint8'))  
        temp_image_path = "temp_image.jpg"
        pil_image.save(temp_image_path)

        x, original_image, pads = load_and_preprocess_image(temp_image_path) 
        mask = segment_cloth(x, model, pads)

        # Ensure mask has same dimensions as original image before saving
        mask = unpad(mask, pads)

        mask_path = "temp_mask.jpg"
        plt.imsave(mask_path, mask, cmap='gray')

        output_image = perform_inpainting(temp_image_path, mask_path, prompt)

        # Resize the output image to match the original image's dimensions
        output_image = resize_and_upscale(output_image, original_image.shape[1], original_image.shape[0])

        return output_image
    except Exception as e:
        raise gr.Error(f"Error processing image: {e}")



with gr.Blocks() as demo:
    with gr.Row():
        # Header with Image and Description
        gr.HTML(
            """
            <div style="display: flex; align-items: center;">
                <img src="https://ant.dpu.ac.th/wp-content/uploads/2024/04/dpulogo.png" style="width: 100px; margin-right: 20px;">
                <div>
                    <h1>Cloth Image Segmentation and Inpainting</h1>
                    <p>This research project explores cloth segmentation using a specialized library, followed by inpainting with Stable Diffusion using a new prompt. It is conducted by the College of Creative Design and Entertainment Technology, Dhurakij Pundit University, in the lab of Asst. Prof. Banyapon Poolsawas under the MIT License.</p>
                </div>
            </div>
            """
        )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Image")
            prompt_input = gr.Textbox(label="Inpainting Prompt (Optional)", value="Chinese Red and Golder Armor")
            run_button = gr.Button("Run")
        with gr.Column():
            image_output = gr.Image(label="Result")

    run_button.click(fn=image_segmentation_and_inpainting, inputs=[image_input, prompt_input], outputs=image_output)

demo.launch(share=True)