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

# بارگذاری مدل
model_path = "sapiens_0.3b_render_people_epoch_100_torchscript.pt2"
model = torch.jit.load(model_path, map_location=torch.device('cpu'))
model.eval()

def predict(image):
    try:
        print("Predict function called")

        # تغییر اندازه تصویر به 768x1024
        image = image.resize((768, 1024))  # تغییر اندازه به 768x1024

        # پیش‌پردازش تصویر
        image = image.convert("RGB")
        input_tensor = np.array(image)
        input_tensor = input_tensor.transpose(2, 0, 1)  # تبدیل از HWC به CHW
        input_tensor = input_tensor[np.newaxis, :]  # افزودن بعد batch
        input_tensor = input_tensor / 255.0  # نرمال‌سازی
        input_tensor = torch.from_numpy(input_tensor).float()

        print(f"Input tensor shape: {input_tensor.shape}")

        # اجرای مدل
        with torch.no_grad():
            output = model(input_tensor)

        print(f"Output tensor shape: {output.shape}")

        # پس‌پردازش خروجی
        output_image = output.squeeze().cpu().numpy()

        # اگر خروجی تک کاناله است، به RGB تبدیل می‌شود
        if output_image.ndim == 2:  # فقط در صورتی که تک کاناله است
            output_image = np.stack([output_image] * 3, axis=-1)
        
        elif output_image.shape[0] == 1:  # اگر کانال اول 1 است، آن را به RGB تبدیل کنید
            output_image = np.tile(output_image, (3, 1, 1))
            output_image = output_image.transpose(1, 2, 0)

        output_image = (output_image * 255).astype(np.uint8)
        output_image = Image.fromarray(output_image)

        print("Output image generated successfully")
        return output_image

    except Exception as e:
        print(f"Error during prediction: {str(e)}")
        return None

# تعریف رابط Gradio
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Input Image"),
    outputs=gr.Image(type="pil", label="Output Image"),
    title="Sapiens Model Inference",
    description="Upload an image to process with the Sapiens model."
)

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