import os
import cv2
import torch
from model import U2NET
from torch.autograd import Variable
import numpy as np
from huggingface_hub import hf_hub_download
import gradio as gr
import math
import ezdxf

# Chuẩn hóa dự đoán
def normPRED(d):
    return (d - torch.min(d)) / (torch.max(d) - torch.min(d))

# Hàm suy luận với U2NET
def inference(net, input_img):
    input_img = input_img / np.max(input_img)
    tmpImg = np.zeros((input_img.shape[0], input_img.shape[1], 3))
    tmpImg[:, :, 0] = (input_img[:, :, 2] - 0.406) / 0.225
    tmpImg[:, :, 1] = (input_img[:, :, 1] - 0.456) / 0.224
    tmpImg[:, :, 2] = (input_img[:, :, 0] - 0.485) / 0.229
    tmpImg = torch.from_numpy(tmpImg.transpose((2, 0, 1))[np.newaxis, :, :, :]).type(torch.FloatTensor)
    tmpImg = Variable(tmpImg.cuda() if torch.cuda.is_available() else tmpImg)
    d1, d2, d3, d4, d5, d6, d7 = net(tmpImg)
    pred = normPRED(1.0 - d1[:, 0, :, :])
    return pred.cpu().data.numpy().squeeze()

# Hàm chính để xử lý ảnh đầu vào và trả về ảnh chân dung và DWF file
def process_image(img, bw_option):
    # Chuyển đổi ảnh thành đen trắng nếu được chọn
    if bw_option:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)  # Chuyển lại thành ảnh 3 kênh cho mô hình
    
    # Chạy suy luận để tạo ảnh chân dung
    result = inference(u2net, img)
    
    # Phát hiện và lấy contours từ ảnh chân dung
    _, threshold = cv2.threshold(np.uint8(result * 255), 0, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(threshold, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    
    # Tạo DWF file từ các contours
    doc = ezdxf.new('R2010')
    msp = doc.modelspace()
    for contour in contours:
        points = [tuple(p[0]) for p in contour]
        msp.add_polyline2d(points)
    
    # Lưu DWF file
    dwf_file = 'portrait_result.dxf'
    doc.saveas(dwf_file)
    
    return (result * 255).astype(np.uint8), dwf_file

# Tải mô hình từ Hugging Face Hub
def load_u2net_model():
    model_path = hf_hub_download(repo_id="Arrcttacsrks/U2net", filename="u2net_portrait.pth", use_auth_token=os.getenv("HF_TOKEN"))
    net = U2NET(3, 1)
    net.load_state_dict(torch.load(model_path, map_location="cuda" if torch.cuda.is_available() else "cpu"))
    net.eval()
    return net

# Khởi tạo mô hình U2NET
u2net = load_u2net_model()

# Tạo giao diện với Gradio
iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="numpy", label="Upload your image"),
        gr.Checkbox(label="Convert to Black & White?", value=False)
    ],
    outputs=[
        gr.Image(type="numpy", label="Portrait Result"),
        gr.File(label="DWF File")
    ],
    title="Portrait Generation with U2NET",
    description="Upload an image to generate its portrait and DWF file."
)
iface.launch()