import gradio as gr
import open_clip
import torch
from PIL import Image

# Load model and tokenizer
model, preprocess = open_clip.create_model_from_pretrained('hf-hub:woweenie/open-clip-vit-h-nsfw-finetune', device='cpu')
tokenizer = open_clip.get_tokenizer('hf-hub:woweenie/open-clip-vit-h-nsfw-finetune')

# Define labels
type_labels = ['2.5d render', '3d render', 'photograph', 'anime drawing', 'drawing', 'illustration', 'painting', 'pre-raphaelite painting', 'concept artwork', 'screenshot']
scene_labels = ['in an airport', 'in the bath', 'on a bed', 'in bed', 'in a bedroom', 'at the beach', 'on a boat', 'in a tent', 'in a car', 'on a chair', 'in the city', 'in a dressing room', 'on the floor', 'at the gym', 'in a hotel room', 'in a kitchen', 'in a living room', 'in an office', 'by a harbor', 'on a bench', 'in a park', 'by a piano', 'on a forest road', 'in a forest', 'in a garden', 'at a lake', 'on the grass', 'on the ground', 'on a paved surface', 'outdoors, on a rock', 'outdoors, on a rug', 'outdoors, on a towel', 'in a photo studio', 'at the pool', 'at a river', 'on a road', 'by the sea', 'showering', 'in the shower', 'on a stool', 'on a rug', 'on a rock', 'on a sofa', 'on a table', 'at a table', 'in a store', 'on snow', 'by a waterfall', 'with a water feature', 'on a windowsill']
expression_labels = ['scared', 'annoyed', 'aroused', 'bored', 'confident', 'distracted', 'dominating', 'embarrassed', 'scared', 'laughing', 'shy', 'orgasm']
clothing_labels = ['a bikini that is too small', 'bikini bottoms', 'a bikini top', 'a bikini', 'a bodysuit', 'a bra', 'a crop top', 'a dress', 'garters', 'glasses', 'goggles', 'gym shorts', 'a halter top', 'a hat', 'a handbra', 'a hoodie', 'a jacket', 'jeans', 'a jumper', 'a gown', 'a lace-up top', 'leggings', 'lingerie', 'a long sleeved top', 'a off-shoulder top', 'a nightgown', 'a coat', 'overalls', 'pink pajamas', 'pajamas', 'panties', 'pantyhose', 'a t-shirt', 'a robe', 'a bathrobe', 'a piece of fabric', 'a scarf', 'a shirt', 'shorts', 'a skirt', 'a sleeveless top', 'a slip', 'sneakers', 'tube socks', 'a sports bra', 'sunglasses', 'sweatpants', 'a one piece swimsuit', 'a t-shirt', 'a tank top', 'a tied shirt', 'a top', 'long pants', 'a wetsuit', 'a backpack', 'high hem', 'see-through', 'short', 'tight','visible nipples']
clothing_labels = ['wearing ' + label for label in clothing_labels]

def process_image(image):
    # Preprocess image
    image = preprocess(image).unsqueeze(0)
    
    # Tokenize labels
    type_text = tokenizer(type_labels)
    scene_text = tokenizer(scene_labels) 
    expression_text = tokenizer(expression_labels)
    clothing_text = tokenizer(clothing_labels)

    with torch.no_grad():
        # Encode image and text
        image_features = model.encode_image(image)
        type_text_features = model.encode_text(type_text)
        scene_text_features = model.encode_text(scene_text)
        expression_text_features = model.encode_text(expression_text)
        clothing_text_features = model.encode_text(clothing_text)
        
        # Normalize features
        image_features /= image_features.norm(dim=-1, keepdim=True)
        type_text_features /= type_text_features.norm(dim=-1, keepdim=True)
        scene_text_features /= scene_text_features.norm(dim=-1, keepdim=True)
        expression_text_features /= expression_text_features.norm(dim=-1, keepdim=True)
        clothing_text_features /= clothing_text_features.norm(dim=-1, keepdim=True)

        # Calculate probabilities
        type_text_probs = (100.0 * image_features @ type_text_features.T).softmax(dim=-1)
        scene_text_probs = (100.0 * image_features @ scene_text_features.T).softmax(dim=-1)
        expression_text_probs = (100.0 * image_features @ expression_text_features.T).softmax(dim=-1)
        clothing_text_probs = (100.0 * image_features @ clothing_text_features.T).softmax(dim=-1)

    # Convert to dictionaries
    type_results = {label: float(type_text_probs[0][i]) for i, label in enumerate(type_labels)}
    scene_results = {label: float(scene_text_probs[0][i]) for i, label in enumerate(scene_labels)}
    expression_results = {label: float(expression_text_probs[0][i]) for i, label in enumerate(expression_labels)}
    clothing_results = {label: float(clothing_text_probs[0][i]) for i, label in enumerate(clothing_labels)}

    return type_results, scene_results, expression_results, clothing_results

# Create Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Label(label="Type Classification", num_top_classes=8),
        gr.Label(label="Scene Classification", num_top_classes=8),
        gr.Label(label="Expression Classification", num_top_classes=8),
        gr.Label(label="Clothing Classification", num_top_classes=8)
    ],
    title="Image Content Moderation",
    description="Upload an image to analyze its content across multiple categories."
)

if __name__ == "__main__":
    iface.launch()