import os
import cv2
import torch  #2.0.1
import numpy as np
import gradio as gr
import torch.nn as nn
import random
import string
from imageAI import SimpleModel
from myImage import ImageToCV,CVtoImage
IMAGE_SIZE = 64

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleModel(path='best_model.pth').to(device)


def predict_and_draw(input_image):
    #input_image=ImageToCV(input_image)
    
    instructions = model.predict(input_image)
    img = cv2.imread(input_image, cv2.IMREAD_GRAYSCALE)
    img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE))
    # 創建一個新的白色圖像
    image = np.zeros((IMAGE_SIZE, IMAGE_SIZE), dtype=np.uint8)
    
    # 執行每條繪圖指令
    for instruction in instructions:
        try:exec(instruction)
        except:
            import traceback
            traceback.print_exc()
    #image=CVtoImage(image)
    return img, image, "\n".join(instructions)


iface = gr.Interface(
    fn=predict_and_draw,
    inputs=gr.Image(type="filepath"),
    outputs=[  gr.Image(type="numpy" ,image_mode="L" ,label="Input Image"), gr.Image(type="numpy" ,image_mode="L" ,label="Output Image"), gr.Textbox(label="Generated Instructions" ,show_copy_button=True)],
    title="Image to Drawing Instructions",
    description="Upload an image, and the model will predict drawing instructions based on it."
)

iface.launch()