Update app.py

app.py

@@ -6,6 +6,7 @@ from torch.autograd import Variable
 import numpy as np
 from huggingface_hub import hf_hub_download
 import gradio as gr
+import ezdxf  # Thêm thư viện ezdxf để tạo file DXF
 
 # Chuẩn hóa dự đoán
 def normPRED(d):
@@ -24,7 +25,22 @@ def inference(net, input_img):
     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
+# Hàm tạo file DXF từ ảnh kết quả
+def convert_to_dxf(image, filename="output.dxf"):
+    # Tìm các đường nét trong ảnh
+    contours, _ = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    
+    # Tạo file DXF và thêm các đường nét
+    doc = ezdxf.new(dxfversion="R2010")
+    msp = doc.modelspace()
+    for contour in contours:
+        points = contour.reshape(-1, 2)
+        msp.add_lwpolyline(points, close=True)
+        
+    doc.saveas(filename)
+    return filename
+
+# Hàm chính để xử lý ảnh đầu vào, trả về ảnh chân dung và lưu DXF
 def process_image(img, bw_option):
     # Chuyển đổi ảnh thành đen trắng nếu được chọn
     if bw_option:
@@ -32,7 +48,12 @@ def process_image(img, bw_option):
         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)
-    return (result * 255).astype(np.uint8)
+    result_img = (result * 255).astype(np.uint8)
+    
+    # Tạo file DXF từ kết quả
+    dxf_path = convert_to_dxf(result_img)
+    
+    return result_img, dxf_path
 
 # Tải mô hình từ Hugging Face Hub
 def load_u2net_model():
@@ -49,12 +70,15 @@ u2net = load_u2net_model()
 iface = gr.Interface(
     fn=process_image,
     inputs=[
-        gr.Image(type="numpy", label="Upload your image"),
-        gr.Checkbox(label="Convert to Black & White?", value=False)  # Thêm tùy chọn tick
+        gr.Image(type="numpy", label="Tải lên ảnh của bạn"),
+        gr.Checkbox(label="Chuyển sang trắng đen?", value=False)
+    ],
+    outputs=[
+        gr.Image(type="numpy", label="Kết quả chân dung"),
+        gr.File(label="Tải xuống file DXF")  # Thêm output cho file DXF
     ],
-    outputs=gr.Image(type="numpy", label="Portrait Result"),
-    title="Portrait Generation with U2NET",
-    description="Upload an image to generate its portrait."
+    title="Tạo ảnh chân dung và file DXF từ ảnh",
+    description="Tải lên một ảnh để tạo ảnh chân dung và file DXF từ ảnh đó."
 )
 
-iface.launch()
+iface.launch()