Update app.py

app.py

@@ -2,93 +2,96 @@ import gradio as gr
 import torch
 from ultralyticsplus import YOLO
 from PIL import Image, ImageDraw, ImageFont
-import os
-from transformers import pipeline
 import numpy as np
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.image import img_to_array
 
 # Tải model YOLO (detection)
-yolo_model_path = "best.pt"  # Thay đường dẫn model YOLO của bạn
+yolo_model_path = "best.pt"
 yolo_model = YOLO(yolo_model_path)
 
-# Tải pipeline classification
-class_pipe = pipeline("image-classification", model="Hemg/Acne-classification")
+# Tải model classification (.h5)
+classification_model_path = "classification.h5"
+class_model = load_model(classification_model_path)
 
-# 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"
+# Danh sách lớp theo thứ tự model classification
+class_labels = ['Acne', 'Comedones', 'Papule', 'Pustule']
+
+# Ánh xạ tiếng Anh -> tiếng Việt
+class_mapping = {
+    'Acne': 'Mụn trứng cá',
+    'Comedones': 'Mụn cám',
+    'Papule': 'Mụn sẩn',
+    'Pustule': 'Mụn mủ'
 }
 
-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
+def detect_and_classify(image, image_size, conf_threshold=0.4, iou_threshold=0.5):
+    # image là đường dẫn file
     pil_image = Image.open(image).convert("RGB")
 
-    # Detect với YOLO
-    results = yolo_model.predict(pil_image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size)
+    # Detect vùng mụn với YOLO
+    results = yolo_model.predict(pil_image, conf=conf_threshold, iou=iou_threshold, imgsz=image_size)
     boxes = results[0].boxes
     num_boxes = len(boxes)
 
-    # 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."
-        return image, f"Tình trạng mụn: {severity}", recommendation, "Không có loại mụn nào được phát hiện."
+        return image, f"Tình trạng mụn: {severity}", recommendation, "Không phát hiện mụn."
 
-    # Chuyển boxes về array
+    # Lấy toạ độ bounding box
     xyxy = boxes.xyxy.detach().cpu().numpy().astype(int)
     confidences = boxes.conf.detach().cpu().numpy()
 
-    # Crop từng box và classify
-    class_names = []
-    for box in xyxy:
-        x1, y1, x2, y2 = box
-        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']  # tiếng Anh
-        class_names.append(top_class)
-
-    # Đánh giá tình trạng dựa trên số lượng mụn
-    if num_boxes > 10:
-        severity = "Nặng"
-        recommendation = "Bạn nên đến gặp bác sĩ da liễu và sử dụng liệu trình trị mụn chuyên sâu."
-    elif 5 <= num_boxes <= 10:
-        severity = "Trung bình"
-        recommendation = "Hãy duy trì skincare đều đặn với sữa rửa mặt dịu nhẹ và dưỡng ẩm phù hợp."
-    else:
-        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 hiện tại."
-
-    # Vẽ bounding box và class name lên ảnh
+    # Chuẩn bị vẽ
     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):
+    # Classification từng box
+    classified_results = []
+    for i, box in enumerate(xyxy, 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)
+        crop = pil_image.crop((x1, y1, x2, y2))
 
-        # Vẽ box
-        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
-        text = f"#{i}: {cname} ({vn_name}) ({conf:.2f})"
+        # Tiền xử lý crop để đưa vào classification model
+        crop_resized = crop.resize((224, 224))
+        img_arr = img_to_array(crop_resized) / 255.0
+        img_arr = np.expand_dims(img_arr, axis=0)
+
+        # Dự đoán loại mụn
+        pred = class_model.predict(img_arr)
+        class_idx = np.argmax(pred)
+        class_name_en = class_labels[class_idx]
+        class_name_vn = class_mapping.get(class_name_en, class_name_en)  # Lấy tên tiếng Việt
+        conf = confidences[i-1]
 
-        # Sử dụng textbbox để lấy kích thước text
+        # Vẽ box và label
+        text = f"#{i}: {class_name_en} ({class_name_vn}) ({conf:.2f})"
         bbox = draw.textbbox((0,0), text, font=font)
         text_w = bbox[2]-bbox[0]
         text_h = bbox[3]-bbox[1]
 
-        # Vẽ nền cho text
+        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
         draw.rectangle([x1, y1 - text_h, x1 + text_w, y1], fill="red")
-        # 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"
+        classified_results.append((i, class_name_en, class_name_vn))
+
+    # Đánh giá tình trạng dựa trên số lượng mụn
+    if num_boxes > 10:
+        severity = "Nặng"
+        recommendation = "Bạn nên đến gặp bác sĩ da liễu và sử dụng liệu trình trị mụn chuyên sâu."
+    elif 5 <= num_boxes <= 10:
+        severity = "Trung bình"
+        recommendation = "Duy trì skincare đều đặn với sữa rửa mặt dịu nhẹ và dưỡng ẩm."
+    else:
+        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."
+
+    # Liệt kê loại mụn
+    acne_types_str = "Danh sách mụn phát hiện:\n"
+    for idx, cname_en, cname_vn in classified_results:
+        acne_types_str += f"Mụn #{idx}: {cname_en} ({cname_vn})\n"
 
     # Lưu ảnh kết quả
     predicted_image_save_path = "predicted_image.jpg"
@@ -97,37 +100,12 @@ def detect_and_classify(image, image_size, conf_thresold=0.4, iou_thresold=0.5):
     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ự, 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ị.
-"""
-
-custom_css = """
-#component-0, #component-1, #component-2, #component-3, #component-4 {
-    font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
-    text-align: center;
-}
-#component-0 h1 {
-    color: #F15B2A;
-}
-#component-0 h2, h3 {
-    color: #333;
-}
-.gr-button {
-    background-color: #F15B2A !important;
-    color: #fff !important;
-    border: none !important;
-    font-weight: bold !important;
-}
-.gr-button:hover {
-    background-color: #d94c1f !important;
-}
+## Ứng dụng Nhận Diện (YOLO) & Phân Loại Mụn (Classification.h5)
+1. Dùng YOLO để phát hiện vùng mụn trên khuôn mặt.
+2. Dùng model classification (h5) để phân loại loại mụn.
+3. Hiển thị kết quả lên ảnh cùng tình trạng da và gợi ý, kèm tên tiếng Anh và tiếng Việt của loại mụn.
 """
 
-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)"),
@@ -142,14 +120,12 @@ outputs = [
     gr.Textbox(label="Loại Mụn Phát Hiện", interactive=False)
 ]
 
-yolo_app = gr.Interface(
+app = gr.Interface(
     fn=detect_and_classify,
     inputs=inputs,
     outputs=outputs,
-    title="YOLOv8 + Classification: Nhận Diện & Phân Loại Mụn",
-    description=description_md,
-    css=custom_css,
-    theme="default"
+    title="YOLO + Classification (H5) Mụn Tiếng Việt",
+    description=description_md
 )
 
-yolo_app.launch(share=True)
+app.launch(share=True)