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) / 2)

            mask = np.zeros_like(image)
            cv2.circle(mask, (center_x, center_y), radius, (255, 255, 255), -1)
            #blurred_image = cv2.GaussianBlur(image, (51, 51), 50)
            blurred_image = cv2.GaussianBlur(image, (201, 201), 200)
            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, parts2blur):
    detector = NudeDetector(model_path="640m.onnx", inference_resolution=640)
    detections = detector.detect(input_img)
    print(detections)

    if len(parts2blur) == 0:
        parts_to_blur = [
        'FEMALE_GENITALIA_EXPOSED', 'MALE_GENITALIA_EXPOSED', 
        'FEMALE_BREAST_EXPOSED', 'BUTTOCKS_EXPOSED',
        'MALE_BREAST_EXPOSED', 'ANUS_EXPOSED'
        ]
    else:
        parts_to_blur = parts2blur
    
    blurred_image_path = circular_blur(input_img, detections, parts_to_blur)
    return blurred_image_path
    
title = "Nudity Detector Vincent version 2.1"
theme = "Nymbo/Alyx_Theme"


with gr.Blocks() as demo:    
        with gr.Row():
            with gr.Column():
                input_image = gr.components.Image(type='filepath')
            with gr.Column():
                output_image = gr.components.Image(type='filepath')    
                
        with gr.Row():        
           submit_btn = gr.Button("submit", scale=1)

        with gr.Row():
            parts2blur = gr.CheckboxGroup(['FEMALE_GENITALIA_EXPOSED', 'MALE_GENITALIA_EXPOSED', 
                'FEMALE_BREAST_EXPOSED', 'BUTTOCKS_EXPOSED',
                'MALE_BREAST_EXPOSED', 'ANUS_EXPOSED'], label="parts", info="Parts to Blur")
            
        gr.Examples(
            examples=[
                [
                    "magazine_sample1.jpeg",
                    #"yolov10x",
                    #640,
                    #0.25,
                    #0.45,
                ],
                [
                    "magazine_sample2.jpeg",
                    #"yolov10m",
                    #640,
                    #0.25,
                    #0.45,
                ],   
                [
                    "demo_sample1.jpg",
                    #"yolov10m",
                    #640,
                    #0.25,
                    #0.45,
                ],      
                [
                    "demo_sample2.jpg",
                    #"yolov10m",
                    #640,
                    #0.25,
                    #0.45,
                ],      
                [
                    "demo_sample3.jpg",
                    #"yolov10m",
                    #640,
                    #0.25,
                    #0.45,
                ],      
                [
                    "demo_sample4.jpg",
                    #"yolov10m",
                    #640,
                    #0.25,
                    #0.45,
                ],      
                [
                    "003071982_0144.jpg",
                    #"yolov10m",
                    #640,
                    #0.25,
                    #0.45,
                ],         
            ],
            fn=process,
            inputs=[
                input_image,
                parts2blur,
                #model_id,
                #image_size,
                #conf_threshold,
                #iou_threshold,
            ],
            outputs=[output_image],
            cache_examples="lazy",
        )

        submit_btn.click(
        fn=process,
        inputs=[input_image, parts2blur],
        outputs=[
            output_image,
        ],
        queue=False,
    )
    

demo.queue().launch(share=True)