ducdatit2002 commited on
Commit
e82b84e
·
verified ·
1 Parent(s): 35b29ce

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +65 -89
app.py CHANGED
@@ -2,93 +2,96 @@ import gradio as gr
2
  import torch
3
  from ultralyticsplus import YOLO
4
  from PIL import Image, ImageDraw, ImageFont
5
- import os
6
- from transformers import pipeline
7
  import numpy as np
 
 
8
 
9
  # Tải model YOLO (detection)
10
- yolo_model_path = "best.pt" # Thay đường dẫn model YOLO của bạn
11
  yolo_model = YOLO(yolo_model_path)
12
 
13
- # Tải pipeline classification
14
- class_pipe = pipeline("image-classification", model="Hemg/Acne-classification")
 
15
 
16
- # Mapping từ tiếng Anh sang tiếng Việt cho các lớp: Comedo, Acne, Clear
17
- label_mapping = {
18
- "Comedo": "Mụn cám",
19
- "Acne": "Mụn trứng cá",
20
- "Clear": "Vùng da sạch"
 
 
 
 
21
  }
22
 
23
- def detect_and_classify(image, image_size, conf_thresold=0.4, iou_thresold=0.5):
24
- # image là đường dẫn file, đọc thành PIL
25
  pil_image = Image.open(image).convert("RGB")
26
 
27
- # Detect với YOLO
28
- results = yolo_model.predict(pil_image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size)
29
  boxes = results[0].boxes
30
  num_boxes = len(boxes)
31
 
32
- # Nếu không có vùng mụn
33
  if num_boxes == 0:
34
  severity = "Tốt"
35
  recommendation = "Làn da bạn khá ổn! Tiếp tục duy trì thói quen chăm sóc da."
36
- return image, f"Tình trạng mụn: {severity}", recommendation, "Không loại mụn nào được phát hiện."
37
 
38
- # Chuyển boxes về array
39
  xyxy = boxes.xyxy.detach().cpu().numpy().astype(int)
40
  confidences = boxes.conf.detach().cpu().numpy()
41
 
42
- # Crop từng box và classify
43
- class_names = []
44
- for box in xyxy:
45
- x1, y1, x2, y2 = box
46
- crop = pil_image.crop((x1, y1, x2, y2))
47
- # Classification trên vùng crop
48
- results_class = class_pipe(crop)
49
- top_class = results_class[0]['label'] # tiếng Anh
50
- class_names.append(top_class)
51
-
52
- # Đánh giá tình trạng dựa trên số lượng mụn
53
- if num_boxes > 10:
54
- severity = "Nặng"
55
- 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."
56
- elif 5 <= num_boxes <= 10:
57
- severity = "Trung bình"
58
- 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."
59
- else:
60
- severity = "Tốt"
61
- 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."
62
-
63
- # Vẽ bounding box và class name lên ảnh
64
  draw = ImageDraw.Draw(pil_image)
65
  font = ImageFont.load_default()
66
 
67
- # Tạo chuỗi kết quả loại mụn
68
- acne_types_str = "Danh sách mụn phát hiện:\n"
69
-
70
- for i, (box, cname, conf) in enumerate(zip(xyxy, class_names, confidences), start=1):
71
  x1, y1, x2, y2 = box
72
-
73
- # Lấy tên tiếng việt từ mapping (nếu không có, dùng tên cũ)
74
- vn_name = label_mapping.get(cname, cname)
75
 
76
- # Vẽ box
77
- draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
78
- text = f"#{i}: {cname} ({vn_name}) ({conf:.2f})"
 
 
 
 
 
 
 
 
79
 
80
- # Sử dụng textbbox để lấy kích thước text
 
81
  bbox = draw.textbbox((0,0), text, font=font)
82
  text_w = bbox[2]-bbox[0]
83
  text_h = bbox[3]-bbox[1]
84
 
85
- # Vẽ nền cho text
86
  draw.rectangle([x1, y1 - text_h, x1 + text_w, y1], fill="red")
87
- # Vẽ text
88
  draw.text((x1, y1 - text_h), text, fill="white", font=font)
89
 
90
- # Thêm vào chuỗi mô tả
91
- acne_types_str += f"Mụn #{i}: {cname} ({vn_name})\n"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
92
 
93
  # Lưu ảnh kết quả
94
  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):
97
  return predicted_image_save_path, f"Tình trạng mụn: {severity}", recommendation, acne_types_str
98
 
99
  description_md = """
100
- ## Ứng dụng Nhận Diện & Phân Loại Mụn
101
- - Sử dụng YOLO để phát hiện các vùng mụn trên khuôn mặt.
102
- - Sử dụng hình phân loại (Hemg/Acne-classification) để xác định loại mụn.
103
- - 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.
104
- - Đánh giá tình trạng da và đưa ra khuyến nghị.
105
- """
106
-
107
- custom_css = """
108
- #component-0, #component-1, #component-2, #component-3, #component-4 {
109
- font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
110
- text-align: center;
111
- }
112
- #component-0 h1 {
113
- color: #F15B2A;
114
- }
115
- #component-0 h2, h3 {
116
- color: #333;
117
- }
118
- .gr-button {
119
- background-color: #F15B2A !important;
120
- color: #fff !important;
121
- border: none !important;
122
- font-weight: bold !important;
123
- }
124
- .gr-button:hover {
125
- background-color: #d94c1f !important;
126
- }
127
  """
128
 
129
- import gradio as gr
130
-
131
  inputs = [
132
  gr.Image(type="filepath", label="Ảnh Khuôn Mặt"),
133
  gr.Slider(minimum=320, maximum=1280, step=32, value=640, label="Kích thước ảnh (Image Size)"),
@@ -142,14 +120,12 @@ outputs = [
142
  gr.Textbox(label="Loại Mụn Phát Hiện", interactive=False)
143
  ]
144
 
145
- yolo_app = gr.Interface(
146
  fn=detect_and_classify,
147
  inputs=inputs,
148
  outputs=outputs,
149
- title="YOLOv8 + Classification: Nhận Diện & Phân Loại Mụn",
150
- description=description_md,
151
- css=custom_css,
152
- theme="default"
153
  )
154
 
155
- yolo_app.launch(share=True)
 
2
  import torch
3
  from ultralyticsplus import YOLO
4
  from PIL import Image, ImageDraw, ImageFont
 
 
5
  import numpy as np
6
+ from tensorflow.keras.models import load_model
7
+ from tensorflow.keras.preprocessing.image import img_to_array
8
 
9
  # Tải model YOLO (detection)
10
+ yolo_model_path = "best.pt"
11
  yolo_model = YOLO(yolo_model_path)
12
 
13
+ # Tải model classification (.h5)
14
+ classification_model_path = "classification.h5"
15
+ class_model = load_model(classification_model_path)
16
 
17
+ # Danh sách lớp theo thứ tự model classification
18
+ class_labels = ['Acne', 'Comedones', 'Papule', 'Pustule']
19
+
20
+ # Ánh xạ tiếng Anh -> tiếng Việt
21
+ class_mapping = {
22
+ 'Acne': 'Mụn trứng cá',
23
+ 'Comedones': 'Mụn cám',
24
+ 'Papule': 'Mụn sẩn',
25
+ 'Pustule': 'Mụn mủ'
26
  }
27
 
28
+ def detect_and_classify(image, image_size, conf_threshold=0.4, iou_threshold=0.5):
29
+ # image là đường dẫn file
30
  pil_image = Image.open(image).convert("RGB")
31
 
32
+ # Detect vùng mụn với YOLO
33
+ results = yolo_model.predict(pil_image, conf=conf_threshold, iou=iou_threshold, imgsz=image_size)
34
  boxes = results[0].boxes
35
  num_boxes = len(boxes)
36
 
 
37
  if num_boxes == 0:
38
  severity = "Tốt"
39
  recommendation = "Làn da bạn khá ổn! Tiếp tục duy trì thói quen chăm sóc da."
40
+ return image, f"Tình trạng mụn: {severity}", recommendation, "Không phát hiện mụn."
41
 
42
+ # Lấy toạ độ bounding box
43
  xyxy = boxes.xyxy.detach().cpu().numpy().astype(int)
44
  confidences = boxes.conf.detach().cpu().numpy()
45
 
46
+ # Chuẩn bị vẽ
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
47
  draw = ImageDraw.Draw(pil_image)
48
  font = ImageFont.load_default()
49
 
50
+ # Classification từng box
51
+ classified_results = []
52
+ for i, box in enumerate(xyxy, start=1):
 
53
  x1, y1, x2, y2 = box
54
+ crop = pil_image.crop((x1, y1, x2, y2))
 
 
55
 
56
+ # Tiền xử lý crop để đưa vào classification model
57
+ crop_resized = crop.resize((224, 224))
58
+ img_arr = img_to_array(crop_resized) / 255.0
59
+ img_arr = np.expand_dims(img_arr, axis=0)
60
+
61
+ # Dự đoán loại mụn
62
+ pred = class_model.predict(img_arr)
63
+ class_idx = np.argmax(pred)
64
+ class_name_en = class_labels[class_idx]
65
+ class_name_vn = class_mapping.get(class_name_en, class_name_en) # Lấy tên tiếng Việt
66
+ conf = confidences[i-1]
67
 
68
+ # Vẽ box label
69
+ text = f"#{i}: {class_name_en} ({class_name_vn}) ({conf:.2f})"
70
  bbox = draw.textbbox((0,0), text, font=font)
71
  text_w = bbox[2]-bbox[0]
72
  text_h = bbox[3]-bbox[1]
73
 
74
+ draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
75
  draw.rectangle([x1, y1 - text_h, x1 + text_w, y1], fill="red")
 
76
  draw.text((x1, y1 - text_h), text, fill="white", font=font)
77
 
78
+ classified_results.append((i, class_name_en, class_name_vn))
79
+
80
+ # Đánh giá tình trạng dựa trên số lượng mụn
81
+ if num_boxes > 10:
82
+ severity = "Nặng"
83
+ 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."
84
+ elif 5 <= num_boxes <= 10:
85
+ severity = "Trung bình"
86
+ recommendation = "Duy trì skincare đều đặn với sữa rửa mặt dịu nhẹ và dưỡng ẩm."
87
+ else:
88
+ severity = "Tốt"
89
+ recommendation = "Làn da bạn khá ổn! Tiếp tục duy trì thói quen chăm sóc da."
90
+
91
+ # Liệt kê loại mụn
92
+ acne_types_str = "Danh sách mụn phát hiện:\n"
93
+ for idx, cname_en, cname_vn in classified_results:
94
+ acne_types_str += f"Mụn #{idx}: {cname_en} ({cname_vn})\n"
95
 
96
  # Lưu ảnh kết quả
97
  predicted_image_save_path = "predicted_image.jpg"
 
100
  return predicted_image_save_path, f"Tình trạng mụn: {severity}", recommendation, acne_types_str
101
 
102
  description_md = """
103
+ ## Ứng dụng Nhận Diện (YOLO) & Phân Loại Mụn (Classification.h5)
104
+ 1. Dùng YOLO để phát hiện vùng mụn trên khuôn mặt.
105
+ 2. Dùng model classification (h5) để phân loại loại mụn.
106
+ 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.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
107
  """
108
 
 
 
109
  inputs = [
110
  gr.Image(type="filepath", label="Ảnh Khuôn Mặt"),
111
  gr.Slider(minimum=320, maximum=1280, step=32, value=640, label="Kích thước ảnh (Image Size)"),
 
120
  gr.Textbox(label="Loại Mụn Phát Hiện", interactive=False)
121
  ]
122
 
123
+ app = gr.Interface(
124
  fn=detect_and_classify,
125
  inputs=inputs,
126
  outputs=outputs,
127
+ title="YOLO + Classification (H5) Mụn Tiếng Việt",
128
+ description=description_md
 
 
129
  )
130
 
131
+ app.launch(share=True)