Update app.py

app.py

@@ -13,6 +13,13 @@ yolo_model = YOLO(yolo_model_path)
 # Tải pipeline classification
 class_pipe = pipeline("image-classification", model="Hemg/Acne-classification")
 
+# Mapping từ tiếng Anh sang tiếng Việt cho các lớp: Comedo, Acne, Clear
+label_mapping = {
+    "Comedo": "Mụn cám",
+    "Acne": "Mụn trứng cá",
+    "Clear": "Vùng da sạch"
+}
+
 def detect_and_classify(image, image_size, conf_thresold=0.4, iou_thresold=0.5):
     # image là đường dẫn file, đọc thành PIL
     pil_image = Image.open(image).convert("RGB")
@@ -22,7 +29,7 @@ def detect_and_classify(image, image_size, conf_thresold=0.4, iou_thresold=0.5):
     boxes = results[0].boxes
     num_boxes = len(boxes)
 
-    # Nếu không có vùng mụn, trả về kết quả
+    # Nếu không có vùng mụn
     if num_boxes == 0:
         severity = "Tốt"
         recommendation = "Làn da bạn khá ổn! Tiếp tục duy trì thói quen chăm sóc da."
@@ -39,7 +46,7 @@ def detect_and_classify(image, image_size, conf_thresold=0.4, iou_thresold=0.5):
         crop = pil_image.crop((x1, y1, x2, y2))
         # Classification trên vùng crop
         results_class = class_pipe(crop)
-        top_class = results_class[0]['label']
+        top_class = results_class[0]['label']  # tiếng Anh
         class_names.append(top_class)
 
     # Đánh giá tình trạng dựa trên số lượng mụn
@@ -57,12 +64,20 @@ def detect_and_classify(image, image_size, conf_thresold=0.4, iou_thresold=0.5):
     draw = ImageDraw.Draw(pil_image)
     font = ImageFont.load_default()
 
+    # Tạo chuỗi kết quả loại mụn
+    acne_types_str = "Danh sách mụn phát hiện:\n"
+
     for i, (box, cname, conf) in enumerate(zip(xyxy, class_names, confidences), start=1):
         x1, y1, x2, y2 = box
+        
+        # Lấy tên tiếng việt từ mapping (nếu không có, dùng tên cũ)
+        vn_name = label_mapping.get(cname, cname)
+
+        # Vẽ box
         draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
-        text = f"#{i}: {cname} ({conf:.2f})"
+        text = f"#{i}: {cname} ({vn_name}) ({conf:.2f})"
 
-        # Dùng textbbox để xác định kích thước text
+        # Sử dụng textbbox để lấy kích thước text
         bbox = draw.textbbox((0,0), text, font=font)
         text_w = bbox[2]-bbox[0]
         text_h = bbox[3]-bbox[1]
@@ -72,22 +87,20 @@ def detect_and_classify(image, image_size, conf_thresold=0.4, iou_thresold=0.5):
         # Vẽ text
         draw.text((x1, y1 - text_h), text, fill="white", font=font)
 
+        # Thêm vào chuỗi mô tả
+        acne_types_str += f"Mụn #{i}: {cname} ({vn_name})\n"
+
     # Lưu ảnh kết quả
     predicted_image_save_path = "predicted_image.jpg"
     pil_image.save(predicted_image_save_path)
 
-    # Liệt kê loại mụn theo số thứ tự
-    acne_types_str = "Danh sách mụn phát hiện:\n"
-    for i, cname in enumerate(class_names, start=1):
-        acne_types_str += f"Mụn #{i}: {cname}\n"
-
     return predicted_image_save_path, f"Tình trạng mụn: {severity}", recommendation, acne_types_str
 
 description_md = """
 ## Ứng dụng Nhận Diện & Phân Loại Mụn
 - Sử dụng YOLO để phát hiện các vùng mụn trên khuôn mặt.
 - Sử dụng mô hình phân loại (Hemg/Acne-classification) để xác định loại mụn.
-- Hiển thị bounding box kèm số thứ tự và loại mụn tương ứng.
+- Hiển thị bounding box kèm số thứ tự, tên tiếng Anh và tiếng Việt của loại mụn.
 - Đánh giá tình trạng da và đưa ra khuyến nghị.
 """
 
@@ -113,6 +126,8 @@ custom_css = """
 }
 """
 
+import gradio as gr
+
 inputs = [
     gr.Image(type="filepath", label="Ảnh Khuôn Mặt"),
     gr.Slider(minimum=320, maximum=1280, step=32, value=640, label="Kích thước ảnh (Image Size)"),