Update data_processing.py

utils/data_processing.py

@@ -76,6 +76,7 @@ def process_image_yolo(image, conf_threshold, iou_threshold, label_mode):
     img = image.copy()
     harmful_label_list = []
     annotations = []
+    label_counts = {}
 
     for i, det in enumerate(pred):
         if len(det):
@@ -83,8 +84,9 @@ def process_image_yolo(image, conf_threshold, iou_threshold, label_mode):
 
             for *xyxy, conf, cls in reversed(det):
                 c = int(cls)
-                if c != 6:
-                    harmful_label_list.append(c)
+                harmful_label_list.append(c)
+                label_name = names[c]
+                label_counts[label_name] = label_counts.get(label_name, 0) + 1
 
                 annotation = {
                     "xyxy": xyxy,
@@ -98,14 +100,7 @@ def process_image_yolo(image, conf_threshold, iou_threshold, label_mode):
                 }
                 annotations.append(annotation)
 
-    if 4 in harmful_label_list and 10 in harmful_label_list:
-        gr.Warning("Warning: This image is featuring underwear.")
-    elif harmful_label_list:
-        gr.Error("Warning: This image may contain harmful content.")
-        img = cv2.GaussianBlur(img, (125, 125), 0)
-    else:
-        gr.Info("This image appears to be safe.")
-
+    # Annotate image
     annotator = Annotator(img, line_width=3, example=str(names))
 
     for ann in annotations:
@@ -122,7 +117,7 @@ def process_image_yolo(image, conf_threshold, iou_threshold, label_mode):
                 -1,
             )
 
-    return annotator.result()
+    return annotator.result(), label_counts
 
 
 def detect_nsfw(input_image, conf_threshold=0.3, iou_threshold=0.45, label_mode="Draw box"):
@@ -145,6 +140,19 @@ def detect_nsfw(input_image, conf_threshold=0.3, iou_threshold=0.45, label_mode=
     #     image_np = resize_and_pad(image_np, imgsz)  # 여기서 imgsz는 (640, 640)
     #     processed_image = process_image_yolo(image_np, conf_threshold, iou_threshold, label_mode)
     #     return "Detailed analysis completed. See the image for results.", processed_image
+
+    # Resize image and process with YOLO
     image_np = resize_and_pad(image_np, imgsz)  # 여기서 imgsz는 (640, 640)
-    processed_image = process_image_yolo(image_np, conf_threshold, iou_threshold, label_mode)
-    return "Detailed analysis completed. See the image for results.", processed_image
+    processed_image, label_counts = process_image_yolo(
+        image_np, conf_threshold, iou_threshold, label_mode
+    )
+
+    # Construct detailed result text
+    if label_counts:
+        result_text = "Detected content:\n"
+        for label, count in label_counts.items():
+            result_text += f" - {label}: {count} instance(s)\n"
+    else:
+        result_text = "This image appears to be safe."
+
+    return result_text, processed_image