Update app.py

app.py

@@ -1,30 +1,48 @@
 import gradio as gr
 import torch
-from ultralyticsplus import YOLO, render_result
-from PIL import Image
+from ultralyticsplus import YOLO
+from PIL import Image, ImageDraw, ImageFont
 import os
+from transformers import pipeline
+import numpy as np
 
-def yolov8_func(image, 
-                image_size, 
-                conf_thresold=0.4,
-                iou_thresold=0.50):
+# Tải model YOLO (detection)
+yolo_model_path = "best.pt"  # Thay đường dẫn model YOLO của bạn
+yolo_model = YOLO(yolo_model_path)
 
-    # Load the YOLOv8 model
-    model_path = "best.pt"
-    model = YOLO(model_path)  # Thay bằng đường dẫn model của bạn
+# Tải pipeline classification
+class_pipe = pipeline("image-classification", model="Hemg/Acne-classification")
 
-    # Dự đoán
-    result = model.predict(image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size)
-    boxes = result[0].boxes
+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")
+
+    # Detect với YOLO
+    results = yolo_model.predict(pil_image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size)
+    boxes = results[0].boxes
     num_boxes = len(boxes)
 
-    # In thông tin (tùy chọn)
-    print("Object type: ", boxes.cls)
-    print("Confidence: ", boxes.conf)
-    print("Coordinates: ", boxes.xyxy)
-    print(f"Số lượng vùng phát hiện: {num_boxes}")
+    # Nếu không có vùng mụn, trả về kết quả
+    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."
+
+    # Chuyển boxes về array
+    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']
+        class_names.append(top_class)
 
-    # Phân loại mức độ mụn và gợi ý
+    # Đá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."
@@ -35,31 +53,44 @@ def yolov8_func(image,
         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."
 
-    print(f"Tình trạng mụn: {severity}")
-    print(f"Khuyến nghị: {recommendation}")
+    # Vẽ bounding box và class name lên ảnh
+    draw = ImageDraw.Draw(pil_image)
+    font = ImageFont.load_default()
 
-    # Kết xuất ảnh kết quả
-    render = render_result(model=model, image=image, result=result[0])
+    for i, (box, cname, conf) in enumerate(zip(xyxy, class_names, confidences), start=1):
+        x1, y1, x2, y2 = box
+        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
+        text = f"#{i}: {cname} ({conf:.2f})"
+
+        # Dùng textbbox để xác định 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]
+
+        # Vẽ nền cho text
+        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)
+
+    # Lưu ảnh kết quả
     predicted_image_save_path = "predicted_image.jpg"
-    render.save(predicted_image_save_path)
+    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
+    return predicted_image_save_path, f"Tình trạng mụn: {severity}", recommendation, acne_types_str
 
-# Tạo phần mô tả và hướng dẫn
 description_md = """
-## Ứng dụng Nhận Diện Mụn bằng YOLOv8
-Chọn một ảnh chụp khuôn mặt của bạn. Ứng dụng sẽ phân tích và đánh giá tình trạng mụn dựa trên mô hình YOLOv8.
-- **Bước 1:** Tải lên ảnh khuôn mặt.
-- **Bước 2:** Điều chỉnh các tham số (nếu cần).
-- **Bước 3:** Nhấn **Submit** để nhận kết quả.
-
-Phần kết quả sẽ cho ra:
-- Ảnh với bounding box quanh các nốt mụn được phát hiện.
-- Đánh giá tình trạng mụn.
-- Lời khuyên chăm sóc da.
+## Ứ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.
+- Đánh giá tình trạng da và đưa ra khuyến nghị.
 """
 
-# CSS tùy chỉnh để giao diện đẹp hơn
 custom_css = """
 #component-0, #component-1, #component-2, #component-3, #component-4 {
     font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
@@ -83,26 +114,27 @@ custom_css = """
 """
 
 inputs = [
-    gr.Image(type="filepath", label="Ảnh Khuôn Mặt (Input Image)"),
+    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)"),
-    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.15, label="Ngưỡng Confidence"),
-    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.15, label="Ngưỡng IOU")
+    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.4, label="Ngưỡng Confidence"),
+    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Ngưỡng IOU")
 ]
 
 outputs = [
     gr.Image(type="filepath", label="Ảnh Sau Khi Xử Lý"),
     gr.Textbox(label="Tình Trạng Mụn", interactive=False),
-    gr.Textbox(label="Khuyến Nghị", interactive=False)
+    gr.Textbox(label="Khuyến Nghị", interactive=False),
+    gr.Textbox(label="Loại Mụn Phát Hiện", interactive=False)
 ]
 
 yolo_app = gr.Interface(
-    fn=yolov8_func,
+    fn=detect_and_classify,
     inputs=inputs,
     outputs=outputs,
-    title="YOLOv8: Nhận Diện Mụn",
+    title="YOLOv8 + Classification: Nhận Diện & Phân Loại Mụn",
     description=description_md,
     css=custom_css,
-    theme="default"  # Bạn có thể thử "default" hoặc "huggingface"
+    theme="default"
 )
 
-yolo_app.launch(debug=True)
+yolo_app.launch(share=True)