Update app.py

app.py

@@ -6,14 +6,12 @@ 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
+import ezdxf
 from PIL import Image, UnidentifiedImageError
 
-# 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))
@@ -22,49 +20,63 @@ 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, _, _, _, _, _, _ = net(tmpImg)
+    d1, *, *, *, *, *, * = net(tmpImg)
     pred = normPRED(1.0 - d1[:, 0, :, :])
     return pred.cpu().data.numpy().squeeze()
 
-# Hàm tạo file DXF từ ảnh kết quả
-def convert_to_dxf(image, filename="output.dxf"):
+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(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    contours, _ = cv2.findContours(portrait_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
     
-    # Tạo file DXF và thêm các đường nét
+    # 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
+    
+    return contours
+
+def convert_to_dxf(contours, filename="output.dxf"):
+    """
+    Tạo file DXF từ các đường nét (contours).
+    
+    Parameters:
+    contours (list): Danh sách các đường nét.
+    filename (str): Tên file DXF.
+    
+    Returns:
+    str: Đường dẫn đến file DXF.
+    """
     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):
     try:
-        # Đảm bảo ảnh đầu vào hợp lệ
-        img = Image.open(img).convert("RGB")  # Chuyển ảnh sang RGB nếu cần
+        img = Image.open(img).convert("RGB")
         img = np.array(img)
-
-        # 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
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
         result = inference(u2net, img)
         result_img = (result * 255).astype(np.uint8)
-        
-        # Tạo file DXF từ kết quả
-        dxf_path = convert_to_dxf(result_img)
-        
+        contours = extract_contours(result_img)
+        dxf_path = convert_to_dxf(contours)
         return result_img, dxf_path
     except UnidentifiedImageError:
-        return "Lỗi: Không thể nhận diện file ảnh. Hãy đảm bảo file đầu vào là ảnh hợp lệ.", None
+        return "Error: Unable to identify the image file. Please ensure the input file is a valid image.", None
 
-# 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)
@@ -72,22 +84,20 @@ def load_u2net_model():
     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="filepath", label="Tải lên ảnh của bạn"),
-        gr.Checkbox(label="Chuyển sang trắng đen?", value=False)
+        gr.Image(type="filepath", label="Upload your image"),
+        gr.Checkbox(label="Convert to black and white?", 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
+        gr.Image(type="numpy", label="Portrait result"),
+        gr.File(label="Download DXF file")
     ],
-    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 đó."
+    title="Create Portrait Images and DXF Files from Images",
+    description="Upload an image to generate a portrait image and a DXF file from it."
 )
 
-iface.launch(share=True)
+iface.launch(share=True)