import gradio as gr
from nudenet import NudeDetector
import cv2
import numpy as np

def circular_blur(image_path, detections, parts_to_blur):
    image = cv2.imread(image_path)
    
    for detection in detections:
        label = detection['class']
        if label in parts_to_blur:
            x, y, width, height = map(int, detection['box'])
            center_x, center_y = x + width // 2, y + height // 2
            radius = int(min(width, height))

            mask = np.zeros_like(image)
            cv2.circle(mask, (center_x, center_y), radius, (255, 255, 255), -1)
            blurred_image = cv2.GaussianBlur(image, (75, 75), 50)
            image = np.where(mask == 255, blurred_image, image)
    
    blurred_image_path = 'blurred_' + image_path.split('/')[-1]
    cv2.imwrite(blurred_image_path, image)
    return blurred_image_path

def process(input_img):
    detector = NudeDetector(model_path="640m.onnx", inference_resolution=640)
    detections = detector.detect(input_img)
    print(detections)
    parts_to_blur = [
        'FEMALE_GENITALIA_EXPOSED', 'MALE_GENITALIA_EXPOSED', 
        'FEMALE_BREAST_EXPOSED', 'BUTTOCKS_EXPOSED',
        'MALE_BREAST_EXPOSED', 'ANUS_EXPOSED'
    ]
    
    blurred_image_path = circular_blur(input_img, detections, parts_to_blur)
    return blurred_image_path
    
title = "SFW Converter"
description = 'Blurs NSFW images using "NudeNet" Package.'

iface = gr.Interface(process, gr.components.Image(type='filepath'), gr.components.Image(type="filepath"), title=title, description=description)

iface.launch()