feat Only detail detect model

app.py

@@ -1,9 +1,10 @@
-import gradio as gr
+import os
 import sys
 from pathlib import Path
-import os
+import gradio as gr
 from utils.data_processing import detect_nsfw
-# Import YOLO-related modules
+
+# YOLO-related module path setup
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[0]
 if str(ROOT) not in sys.path:
@@ -11,37 +12,43 @@ if str(ROOT) not in sys.path:
 ROOT = Path(os.path.relpath(ROOT, Path.cwd()))
 
 
-
-
-
-# Gradio interface
+# Define the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# NSFW Content Detection")
-    with gr.Row():
-        detection_mode = gr.Radio(["Simple Check", "Detailed Analysis"], label="Detection Mode", value="Simple Check")
+    gr.Markdown("# NSFW Content Detection - Detailed Analysis")
+
+    # Advanced parameters for Detailed Analysis
     with gr.Row():
-        conf_threshold = gr.Slider(0, 1, value=0.3, label="Confidence Threshold", visible=False)
-        iou_threshold = gr.Slider(0, 1, value=0.45, label="Overlap Threshold", visible=False)
-        label_mode = gr.Dropdown(["Draw box", "Draw Label", "Draw Confidence", "Censor Predictions"], label="Label Display Mode", value="Draw box", visible=False)
-    
+        conf_threshold = gr.Slider(0, 1, value=0.3, label="Confidence Threshold")
+        iou_threshold = gr.Slider(0, 1, value=0.45, label="Overlap Threshold")
+        label_mode = gr.Dropdown(
+            ["Draw box", "Draw Label", "Draw Confidence", "Censor Predictions"],
+            label="Label Display Mode",
+            value="Draw box",
+        )
+
+    # Input and output components
     with gr.Row():
         input_image = gr.Image(type="numpy", label="Upload an image or enter a URL")
         output_text = gr.Textbox(label="Detection Result")
     with gr.Row():
-        output_image = gr.Image(type="numpy", label="Processed Image (for detailed analysis)", visible=False)
+        output_image = gr.Image(type="numpy", label="Processed Image (for detailed analysis)")
 
+    # Detection button
     detect_button = gr.Button("Detect")
-    
-    def update_visibility(mode):
-        return [gr.update(visible=(mode == "Detailed Analysis"))] * 4
-    
-    detection_mode.change(update_visibility, inputs=[detection_mode], outputs=[conf_threshold, iou_threshold, label_mode, output_image])
-    
+
+    # Connect detection button to the detect_nsfw function
+    def safe_detect_nsfw(image, conf, iou, label):
+        try:
+            return detect_nsfw(image, "Detailed Analysis", conf, iou, label)
+        except Exception as e:
+            return f"Error during detection: {e}", None
+
     detect_button.click(
-        detect_nsfw,
-        inputs=[input_image, detection_mode, conf_threshold, iou_threshold, label_mode],
-        outputs=[output_text, output_image]
+        safe_detect_nsfw,
+        inputs=[input_image, conf_threshold, iou_threshold, label_mode],
+        outputs=[output_text, output_image],
     )
 
+# Launch the Gradio app
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0")
+    demo.launch(server_name="0.0.0.0")

config/settings.py

@@ -1,7 +1,8 @@
 import os
+
 # 프로젝트 루트 디렉토리 경로 얻기
 BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
 # MODEL_PATH를 절대 경로로 설정
-DETECT_MODEL_PATH = os.path.join(BASE_DIR, 'weights', 'yolov9_c_nsfw.pt')
-CLASSIFICATION_MODEL_PATH = "Falconsai/nsfw_image_detection"
+DETECT_MODEL_PATH = os.path.join(BASE_DIR, "weights", "yolov9_c_nsfw.pt")
+# CLASSIFICATION_MODEL_PATH = "Falconsai/nsfw_image_detection"

models/common.py

@@ -30,7 +30,7 @@ from utils.general import (LOGGER, ROOT, Profile, check_requirements, check_suff
                            xywh2xyxy, xyxy2xywh, yaml_load)
 from utils.plots import Annotator, colors, save_one_box
 from utils.torch_utils import copy_attr, smart_inference_mode
-from config.settings import CLASSIFICATION_MODEL_PATH
+# from config.settings import CLASSIFICATION_MODEL_PATH
 
 def autopad(k, p=None, d=1):  # kernel, padding, dilation
     # Pad to 'same' shape outputs
@@ -1200,17 +1200,17 @@ class Classify(nn.Module):
             x = torch.cat(x, 1)
         return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))
 
-class NSFWModel:
-    def __init__(self):
-        self.model = AutoModelForImageClassification.from_pretrained(CLASSIFICATION_MODEL_PATH)
-        self.processor = ViTImageProcessor.from_pretrained(CLASSIFICATION_MODEL_PATH)
-        self.id2label = self.model.config.id2label
-
-    def predict(self, image: Image.Image) -> str:
-        with torch.no_grad():
-            inputs = self.processor(images=image, return_tensors="pt")
-            outputs = self.model(**inputs)
-            logits = outputs.logits
-            predicted_label = logits.argmax(-1).item()
+# class NSFWModel:
+#     def __init__(self):
+#         self.model = AutoModelForImageClassification.from_pretrained(CLASSIFICATION_MODEL_PATH, local_files_only=True)
+#         self.processor = ViTImageProcessor.from_pretrained(CLASSIFICATION_MODEL_PATH, local_files_only=True)
+#         self.id2label = self.model.config.id2label
+
+#     def predict(self, image: Image.Image) -> str:
+#         with torch.no_grad():
+#             inputs = self.processor(images=image, return_tensors="pt")
+#             outputs = self.model(**inputs)
+#             logits = outputs.logits
+#             predicted_label = logits.argmax(-1).item()
             
-            return self.id2label[predicted_label]
+#             return self.id2label[predicted_label]

utils/data_processing.py

@@ -6,43 +6,53 @@ from io import BytesIO
 import torch
 import gradio as gr
 
-from models.common import DetectMultiBackend, NSFWModel
+from models.common import DetectMultiBackend  # , NSFWModel
 from utils.torch_utils import select_device
-from utils.general import (check_img_size, non_max_suppression, scale_boxes)
+from utils.general import check_img_size, non_max_suppression, scale_boxes
 from utils.plots import Annotator, colors
 from config.settings import DETECT_MODEL_PATH
 
-# Load classification model
-nsfw_model = NSFWModel()
+# # Load classification model
+# nsfw_model = NSFWModel()
 
 # Load YOLO model
-device = select_device('')
+device = select_device("")
 yolo_model = DetectMultiBackend(DETECT_MODEL_PATH, device=device, dnn=False, data=None, fp16=False)
 stride, names, pt = yolo_model.stride, yolo_model.names, yolo_model.pt
 imgsz = check_img_size((640, 640), s=stride)
 
+
 def resize_and_pad(image, target_size):
     ih, iw = image.shape[:2]
     target_h, target_w = target_size
 
     # 이미지의 가로세로 비율 계산
-    scale = min(target_h/ih, target_w/iw)
-    
+    scale = min(target_h / ih, target_w / iw)
+
     # 새로운 크기 계산
     new_h, new_w = int(ih * scale), int(iw * scale)
-    
+
     # 이미지 리사이즈
     resized = cv2.resize(image, (new_w, new_h))
-    
+
     # 패딩 계산
     pad_h = (target_h - new_h) // 2
     pad_w = (target_w - new_w) // 2
-    
+
     # 패딩 추가
-    padded = cv2.copyMakeBorder(resized, pad_h, target_h-new_h-pad_h, pad_w, target_w-new_w-pad_w, cv2.BORDER_CONSTANT, value=[0,0,0])
-    
+    padded = cv2.copyMakeBorder(
+        resized,
+        pad_h,
+        target_h - new_h - pad_h,
+        pad_w,
+        target_w - new_w - pad_w,
+        cv2.BORDER_CONSTANT,
+        value=[0, 0, 0],
+    )
+
     return padded
 
+
 def process_image_yolo(image, conf_threshold, iou_threshold, label_mode):
     # Image preprocessing
     im = torch.from_numpy(image).to(device).permute(2, 0, 1)
@@ -50,15 +60,15 @@ def process_image_yolo(image, conf_threshold, iou_threshold, label_mode):
     im /= 255
     if len(im.shape) == 3:
         im = im[None]
-    
+
     # Resize image
-    im = torch.nn.functional.interpolate(im, size=imgsz, mode='bilinear', align_corners=False)
-    
+    im = torch.nn.functional.interpolate(im, size=imgsz, mode="bilinear", align_corners=False)
+
     # Inference
     pred = yolo_model(im, augment=False, visualize=False)
     if isinstance(pred, list):
         pred = pred[0]
-        
+
     # NMS
     pred = non_max_suppression(pred, conf_threshold, iou_threshold, None, False, max_det=1000)
 
@@ -66,61 +76,75 @@ def process_image_yolo(image, conf_threshold, iou_threshold, label_mode):
     img = image.copy()
     harmful_label_list = []
     annotations = []
-    
+
     for i, det in enumerate(pred):
         if len(det):
             det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], img.shape).round()
-            
+
             for *xyxy, conf, cls in reversed(det):
                 c = int(cls)
                 if c != 6:
                     harmful_label_list.append(c)
-                
+
                 annotation = {
-                    'xyxy': xyxy,
-                    'conf': conf,
-                    'cls': c,
-                    'label': f"{names[c]} {conf:.2f}" if label_mode == "Draw Confidence" else f"{names[c]}"
+                    "xyxy": xyxy,
+                    "conf": conf,
+                    "cls": c,
+                    "label": (
+                        f"{names[c]} {conf:.2f}"
+                        if label_mode == "Draw Confidence"
+                        else f"{names[c]}"
+                    ),
                 }
                 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.')
-    
+        gr.Info("This image appears to be safe.")
+
     annotator = Annotator(img, line_width=3, example=str(names))
-    
+
     for ann in annotations:
         if label_mode == "Draw box":
-            annotator.box_label(ann['xyxy'], None, color=colors(ann['cls'], True))
+            annotator.box_label(ann["xyxy"], None, color=colors(ann["cls"], True))
         elif label_mode in ["Draw Label", "Draw Confidence"]:
-            annotator.box_label(ann['xyxy'], ann['label'], color=colors(ann['cls'], True))
+            annotator.box_label(ann["xyxy"], ann["label"], color=colors(ann["cls"], True))
         elif label_mode == "Censor Predictions":
-            cv2.rectangle(img, (int(ann['xyxy'][0]), int(ann['xyxy'][1])), (int(ann['xyxy'][2]), int(ann['xyxy'][3])), (0, 0, 0), -1)
+            cv2.rectangle(
+                img,
+                (int(ann["xyxy"][0]), int(ann["xyxy"][1])),
+                (int(ann["xyxy"][2]), int(ann["xyxy"][3])),
+                (0, 0, 0),
+                -1,
+            )
 
     return annotator.result()
 
-def detect_nsfw(input_image, detection_mode, conf_threshold=0.3, iou_threshold=0.45, label_mode="Draw box"):
+
+def detect_nsfw(input_image, conf_threshold=0.3, iou_threshold=0.45, label_mode="Draw box"):
     if isinstance(input_image, str):  # URL input
         response = requests.get(input_image)
         image = Image.open(BytesIO(response.content))
     else:  # File upload
         image = Image.fromarray(input_image)
-    
+
     image_np = np.array(image)
     if len(image_np.shape) == 2:  # grayscale
         image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
     elif image_np.shape[2] == 4:  # RGBA
         image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
-    
-    if detection_mode == "Simple Check":
-        result = nsfw_model.predict(image)
-        return result, None
-    else:  # Detailed Analysis
-        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
+
+    # if detection_mode == "Simple Check":
+    #     result = nsfw_model.predict(image)
+    #     return result, None
+    # else:  # Detailed Analysis
+    #     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
+    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