FEAT Fastapi

app.py

@@ -1,126 +1,16 @@
-import gradio as gr
-import cv2
-import numpy as np
-from PIL import Image
-import requests
-from io import BytesIO
-import torch
-import sys
-from pathlib import Path
-import os
-FILE = Path(__file__).resolve()
-ROOT = FILE.parents[0]  # YOLOv5 root directory
-if str(ROOT) not in sys.path:
-    sys.path.append(str(ROOT))  # add ROOT to PATH
-ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
-
-from models.common import DetectMultiBackend
-from utils.general import (check_img_size, non_max_suppression, scale_boxes)
-from utils.plots import Annotator, colors
-from utils.torch_utils import select_device
-
-# YOLOv9 모델 로드
-device = select_device('')
-model = DetectMultiBackend('./weights/nsfw_detector_e_rok.pt', device=device, dnn=False, data=None, fp16=False)
-stride, names, pt = model.stride, model.names, model.pt
-imgsz = check_img_size((640, 640), s=stride)  # check image size
-
-def process_image(image, conf_threshold, iou_threshold, label_mode):
-    # 이미지 전처리
-    im = torch.from_numpy(image).to(device).permute(2, 0, 1)  # HWC to CHW
-    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
-    im /= 255  # 0 - 255 to 0.0 - 1.0
-    if len(im.shape) == 3:
-        im = im[None]  # expand for batch dim
-    
-    # 이미지 크기 조정
-    im = torch.nn.functional.interpolate(im, size=imgsz, mode='bilinear', align_corners=False)
-    
-    # 추론
-    pred = 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)
-
-    # 결과 처리
-    img = image.copy()
-    
-    harmful_label_list = []
-    annotations = []
-    
-    for i, det in enumerate(pred):  # per image
-        if len(det):
-            # Rescale boxes from img_size to im0 size
-            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], img.shape).round()
-            
-            # Write results
-            for *xyxy, conf, cls in reversed(det):
-                c = int(cls)  # integer class
-                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]}"
-                }
-                annotations.append(annotation)
-    
-    if harmful_label_list:
-        gr.Error("Warning, this is a harmful image.")
-        # 이미지 전체를 흐리게 처리
-        img = cv2.GaussianBlur(img, (125, 125), 0)
-    else:
-        gr.Info('This is a safe image.')
-    
-    # Annotator 적용
-    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))
-        elif label_mode in ["Draw Label", "Draw Confidence"]:
-            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)
-
-    return annotator.result()
-
-def detect_nsfw(input_image, conf_threshold, iou_threshold, label_mode):
-    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.array(image)
-    if len(image.shape) == 2:  # grayscale
-        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-    elif image.shape[2] == 4:  # RGBA
-        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
-    
-    # 이미지 크기 조정
-    image = cv2.resize(image, imgsz)
-    
-    processed_image = process_image(image, conf_threshold, iou_threshold, label_mode)
-    return processed_image
-
-# Gradio 인터페이스 설정
-demo = gr.Interface(
-    fn=detect_nsfw,
-    inputs=[
-        gr.Image(type="numpy", label="Upload an image or enter a URL"),
-        gr.Slider(0, 1, value=0.45, label="Confidence Threshold"),
-        gr.Slider(0, 1, value=0.45, label="Overlap Threshold"),
-        gr.Dropdown(["Draw box", "Draw Label", "Draw Confidence", "Censor Predictions"], label="Label Display Mode", value="Draw box")
-    ],
-    outputs=gr.Image(type="numpy", label="Processed Image"),
-    title="YOLOv9 NSFW Content Detection",
-    description="Upload an image or enter a URL to detect NSFW content using YOLOv9."
-)
-
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0")
+from fastapi import FastAPI, File, UploadFile, Form
+from utils.api import process_image_api
+
+app = FastAPI()
+
+@app.post("/process_image/")
+async def process_image_endpoint(
+    file: UploadFile = File(...),
+    conf_threshold: float = Form(0.25),
+    iou_threshold: float = Form(0.45),
+    label_mode: str = Form("Draw Confidence")
+):
+    return await process_image_api(file, conf_threshold, iou_threshold, label_mode)
+
+# Gunicorn용 app 변수
+application = app

gradio_app.py

@@ -0,0 +1,126 @@
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+import requests
+from io import BytesIO
+import torch
+import sys
+from pathlib import Path
+import os
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[0]  # YOLOv5 root directory
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))  # add ROOT to PATH
+ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+
+from models.common import DetectMultiBackend
+from utils.general import (check_img_size, non_max_suppression, scale_boxes)
+from utils.plots import Annotator, colors
+from utils.torch_utils import select_device
+
+# YOLOv9 모델 로드
+device = select_device('')
+model = DetectMultiBackend('./weights/nsfw_detector_e_rok.pt', device=device, dnn=False, data=None, fp16=False)
+stride, names, pt = model.stride, model.names, model.pt
+imgsz = check_img_size((640, 640), s=stride)  # check image size
+
+def process_image(image, conf_threshold, iou_threshold, label_mode):
+    # 이미지 전처리
+    im = torch.from_numpy(image).to(device).permute(2, 0, 1)  # HWC to CHW
+    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
+    im /= 255  # 0 - 255 to 0.0 - 1.0
+    if len(im.shape) == 3:
+        im = im[None]  # expand for batch dim
+    
+    # 이미지 크기 조정
+    im = torch.nn.functional.interpolate(im, size=imgsz, mode='bilinear', align_corners=False)
+    
+    # 추론
+    pred = 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)
+
+    # 결과 처리
+    img = image.copy()
+    
+    harmful_label_list = []
+    annotations = []
+    
+    for i, det in enumerate(pred):  # per image
+        if len(det):
+            # Rescale boxes from img_size to im0 size
+            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], img.shape).round()
+            
+            # Write results
+            for *xyxy, conf, cls in reversed(det):
+                c = int(cls)  # integer class
+                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]}"
+                }
+                annotations.append(annotation)
+    
+    if harmful_label_list:
+        gr.Error("Warning, this is a harmful image.")
+        # 이미지 전체를 흐리게 처리
+        img = cv2.GaussianBlur(img, (125, 125), 0)
+    else:
+        gr.Info('This is a safe image.')
+    
+    # Annotator 적용
+    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))
+        elif label_mode in ["Draw Label", "Draw Confidence"]:
+            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)
+
+    return annotator.result()
+
+def detect_nsfw(input_image, conf_threshold, iou_threshold, label_mode):
+    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.array(image)
+    if len(image.shape) == 2:  # grayscale
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+    elif image.shape[2] == 4:  # RGBA
+        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+    
+    # 이미지 크기 조정
+    image = cv2.resize(image, imgsz)
+    
+    processed_image = process_image(image, conf_threshold, iou_threshold, label_mode)
+    return processed_image
+
+# Gradio 인터페이스 설정
+demo = gr.Interface(
+    fn=detect_nsfw,
+    inputs=[
+        gr.Image(type="numpy", label="Upload an image or enter a URL"),
+        gr.Slider(0, 1, value=0.45, label="Confidence Threshold"),
+        gr.Slider(0, 1, value=0.45, label="Overlap Threshold"),
+        gr.Dropdown(["Draw box", "Draw Label", "Draw Confidence", "Censor Predictions"], label="Label Display Mode", value="Draw box")
+    ],
+    outputs=gr.Image(type="numpy", label="Processed Image"),
+    title="YOLOv9 NSFW Content Detection",
+    description="Upload an image or enter a URL to detect NSFW content using YOLOv9."
+)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0")

utils/api.py

@@ -0,0 +1,82 @@
+from fastapi import File, UploadFile, Form
+from fastapi.responses import JSONResponse
+from PIL import Image
+import io
+import numpy as np
+import torch
+
+from .general import non_max_suppression, scale_boxes
+from models.common import DetectMultiBackend
+from .torch_utils import select_device
+
+# 모델 로드
+device = select_device('')
+model = DetectMultiBackend('weights/nsfw_detector_e_rok.pt', device=device, dnn=False, data='data/coco128.yaml', fp16=False)
+names = model.names
+imgsz = (640, 640)
+
+async def process_image_api(
+    file: UploadFile = File(...),
+    conf_threshold: float = Form(0.25),
+    iou_threshold: float = Form(0.45),
+    label_mode: str = Form("Draw Confidence")
+):
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents))
+    image_np = np.array(image)
+    
+    result = process_image(image_np, conf_threshold, iou_threshold, label_mode)
+    
+    return JSONResponse(content={"result": result.result})
+
+def process_image(image, conf_threshold, iou_threshold, label_mode):
+    # 이미지 전처리
+    im = torch.from_numpy(image).to(device).permute(2, 0, 1)  # HWC to CHW
+    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
+    im /= 255  # 0 - 255 to 0.0 - 1.0
+    if len(im.shape) == 3:
+        im = im[None]  # expand for batch dim
+    
+    # 이미지 크기 조정
+    im = torch.nn.functional.interpolate(im, size=imgsz, mode='bilinear', align_corners=False)
+    
+    # 추론
+    pred = 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)
+
+    # 결과 처리
+    img = image.copy()
+    
+    harmful_label_list = []
+    annotations = []
+    
+    for i, det in enumerate(pred):  # per image
+        if len(det):
+            # Rescale boxes from img_size to im0 size
+            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], img.shape).round()
+            
+            # Write results
+            for *xyxy, conf, cls in reversed(det):
+                c = int(cls)  # integer class
+                if c != 6:
+                    harmful_label_list.append(c)
+                
+                annotation = {
+                    'xyxy': xyxy,
+                    'conf': float(conf),
+                    'cls': c,
+                    'label': f"{names[c]} {conf:.2f}" if label_mode == "Draw Confidence" else f"{names[c]}"
+                }
+                annotations.append(annotation)
+    
+    result = 'nsfw' if harmful_label_list else 'nomal'
+    return ProcessResponse(result=result)
+
+class ProcessResponse:
+    def __init__(self, result: int):
+        self.result = result
+        # self.annotations = annotations

utils/gunicorn_config.py

@@ -0,0 +1,3 @@
+workers = 1
+worker_class = "uvicorn.workers.UvicornWorker"
+bind = '0.0.0.0:8001'