import gradio as gr
from nudenet import NudeDetector
import cv2

def blur_image(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'])
            x1, y1 = x, y
            x2, y2 = x + width, y + height

            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(image.shape[1], x2)
            y2 = min(image.shape[0], y2)

            print(f"Blurring region: {x1, y1, x2, y2} for label: {label}")
            roi = image[y1:y2, x1:x2]
            
            if roi.size == 0:
                print(f"Skipping empty ROI for detection: {detection}")
                continue
            
            roi = cv2.GaussianBlur(roi, (51, 51), 30)
            image[y1:y2, x1:x2] = roi
    
    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 = blur_image(input_img, detections, parts_to_blur)
    return blurred_image_path

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

iface.launch()