Update app.py

app.py

@@ -6,13 +6,14 @@ from torch.autograd import Variable
 import numpy as np
 from huggingface_hub import hf_hub_download
 import gradio as gr
+import math
 import ezdxf
-from PIL import Image, UnidentifiedImageError
-import xml.etree.ElementTree as ET
 
+# 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))
@@ -21,115 +22,38 @@ def inference(net, input_img):
     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)
+    d1, *, *, *, *, *, * = net(tmpImg)
     pred = normPRED(1.0 - d1[:, 0, :, :])
     return pred.cpu().data.numpy().squeeze()
 
-def extract_contours(portrait_mask):
-    """
-    Trích xuất các đường nét (contours) từ ảnh chân dung.
-    
-    Parameters:
-    portrait_mask (numpy.ndarray): Ảnh chân dung dạng binary (đen trắng).
-    
-    Returns:
-    list: Danh sách các đường nét (contours) được trích xuất.
-    """
-    # Tìm các đường nét trong ảnh
-    contours, _ = cv2.findContours(portrait_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
-    
-    # Làm sạch các đường nét
-    contours = [cnt for cnt in contours if cv2.contourArea(cnt) > 100]  # Loại bỏ các đường nét nhỏ
-    contours = [cv2.approxPolyDP(cnt, 0.01 * cv2.arcLength(cnt, True), True) for cnt in contours]  # Làm trơn đường nét
-    
-    print(f"Number of contours: {len(contours)}")
-    for i, contour in enumerate(contours):
-        print(f"Contour {i}: {contour.shape}")
+# 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
     
-    return contours
-
-def convert_to_dxf(contours, filename="output.dxf"):
-    """
-    Tạo file DXF từ các đường nét (contours).
+    # Chạy suy luận để tạo ảnh chân dung
+    result = inference(u2net, img)
     
-    Parameters:
-    contours (list): Danh sách các đường nét.
-    filename (str): Tên file DXF.
+    # 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)
     
-    Returns:
-    str: Đường dẫn đến file DXF.
-    """
-    doc = ezdxf.new(dxfversion="R2010")
+    # Tạo DWF file từ các contours
+    doc = ezdxf.new('R2010')
     msp = doc.modelspace()
-    
-    # Thêm tất cả các đường nét vào file DXF
     for contour in contours:
-        points = contour.reshape(-1, 2)
-        msp.add_lwpolyline(points, close=True)
-    
-    doc.saveas(filename)
-    return filename
-
-def convert_to_svg(contours, filename="output.svg"):
-    """
-    Tạo file SVG từ các đường nét (contours).
-    
-    Parameters:
-    contours (list): Danh sách các đường nét.
-    filename (str): Tên file SVG.
+        points = [tuple(p[0]) for p in contour]
+        msp.add_polyline2d(points)
     
-    Returns:
-    str: Đường dẫn đến file SVG.
-    """
-    root = ET.Element("svg", {
-        "xmlns": "http://www.w3.org/2000/svg",
-        "width": "100%",
-        "height": "100%",
-        "viewBox": "0 0 100 100"
-    })
-    
-    for contour in contours:
-        if len(contour) > 1:
-            path_data = "M"
-            for point in contour:
-                path_data += f" {point[0]/100} {point[1]/100}"
-            path_data += " Z"
-            path = ET.SubElement(root, "path", {
-                "d": path_data,
-                "fill": "none",
-                "stroke": "black",
-                "stroke-width": "0.5"
-            })
-        elif len(contour) == 1:
-            # Handle the case where the contour has only one point
-            point = contour[0]
-            circle = ET.SubElement(root, "circle", {
-                "cx": f"{point[0][0]/100}",
-                "cy": f"{point[0][1]/100}",
-                "r": "0.5",
-                "fill": "black"
-            })
+    # Lưu DWF file
+    dwf_file = 'portrait_result.dxf'
+    doc.saveas(dwf_file)
     
-    tree = ET.ElementTree(root)
-    tree.write(filename)
-    return filename
-
-def process_image(img, bw_option):
-    try:
-        img = Image.open(img).convert("RGB")
-        img = np.array(img)
-        if bw_option:
-            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-        result = inference(u2net, img)
-        result_img = (result * 255).astype(np.uint8)
-        contours = extract_contours(result_img)
-        dxf_path = convert_to_dxf(contours, filename="output.dxf")
-        svg_path = convert_to_svg(contours, filename="output.svg")
-        return result_img, dxf_path, svg_path
-    except UnidentifiedImageError:
-        return "Error: Unable to identify the image file. Please ensure the input file is a valid image.", None, None
+    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)
@@ -137,25 +61,21 @@ def load_u2net_model():
     net.eval()
     return net
 
+# Khởi tạo mô hình U2NET
 u2net = load_u2net_model()
 
-def main():
-    iface = gr.Interface(
-        fn=process_image,
-        inputs=[
-            gr.Image(type="filepath", label="Upload your image"),
-            gr.Checkbox(label="Convert to black and white?", value=False)
-        ],
-        outputs=[
-            gr.Image(type="numpy", label="Portrait result"),
-            gr.File(label="Download DXF file"),
-            gr.File(label="Download SVG file")
-        ],
-        title="Create Portrait Images, DXF and SVG Files from Images",
-        description="Upload an image to generate a portrait image, a DXF file, and an SVG file from it."
-    )
-
-    iface.launch(share=True)
-
-if __name__ == "__main__":
-    main()
+# 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()