Create app.py

app.py

@@ -0,0 +1,71 @@
+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(), torch.cuda.amp.autocast():
+        # 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=3),
+        gr.Label(label="Scene Classification", num_top_classes=3),
+        gr.Label(label="Expression Classification", num_top_classes=3),
+        gr.Label(label="Clothing Classification", num_top_classes=3)
+    ],
+    title="Image Content Moderation",
+    description="Upload an image to analyze its content across multiple categories."
+)
+
+if __name__ == "__main__":
+    iface.launch()