Update app.py

app.py

@@ -1,28 +1,28 @@
 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')
+DEVICE='cpu'
+model, preprocess = open_clip.create_model_from_pretrained('hf-hub:woweenie/open-clip-vit-h-nsfw-finetune', device=DEVICE)
 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 = ['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']
 clothing_labels = ['wearing ' + label for label in clothing_labels]
 
-def process_image(image):
+def process_image_separate_tags(image):
     # Preprocess image
-    image = preprocess(image).unsqueeze(0)
-    
+    image = preprocess(image).unsqueeze(0).to(DEVICE)
+
     # Tokenize labels
-    type_text = tokenizer(type_labels)
-    scene_text = tokenizer(scene_labels) 
-    expression_text = tokenizer(expression_labels)
-    clothing_text = tokenizer(clothing_labels)
+    type_text = tokenizer(type_labels).to(DEVICE)
+    scene_text = tokenizer(scene_labels).to(DEVICE) 
+    expression_text = tokenizer(expression_labels).to(DEVICE)
+    clothing_text = tokenizer(clothing_labels).to(DEVICE)
 
     with torch.no_grad():
         # Encode image and text
@@ -31,7 +31,7 @@ def process_image(image):
         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)
@@ -39,11 +39,13 @@ def process_image(image):
         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).sigmoid(dim=-1)
-        scene_text_probs = (100.0 * image_features @ scene_text_features.T).sigmoid(dim=-1)
-        expression_text_probs = (100.0 * image_features @ expression_text_features.T).sigmoid(dim=-1)
-        clothing_text_probs = (100.0 * image_features @ clothing_text_features.T).sigmoid(dim=-1)
+        # Calculate cosine similarities and apply softmax
+        # Using temperature parameter to control the "sharpness" of the distribution
+        temperature = 0.1  # Lower values make the distribution more peaked
+        type_text_probs = torch.softmax(image_features @ type_text_features.T / temperature, dim=-1)
+        scene_text_probs = torch.softmax(image_features @ scene_text_features.T / temperature, dim=-1)
+        expression_text_probs = torch.softmax(image_features @ expression_text_features.T / temperature, dim=-1)
+        clothing_text_probs = torch.softmax(image_features @ clothing_text_features.T / temperature, dim=-1)
 
     # Convert to dictionaries
     type_results = {label: float(type_text_probs[0][i]) for i, label in enumerate(type_labels)}
@@ -53,19 +55,68 @@ def process_image(image):
 
     return type_results, scene_results, expression_results, clothing_results
 
+def process_image_combined_tags(image):
+    # Preprocess image
+    image = preprocess(image).unsqueeze(0).to(DEVICE)
+
+    # Tokenize labels
+    all_text = tokenizer(type_labels + scene_labels + expression_labels + clothing_labels).to(DEVICE)
+
+    with torch.no_grad():
+        # Encode image and text
+        image_features = model.encode_image(image)
+        all_text_features = model.encode_text(all_text)
+
+        # Normalize features
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        all_text_features /= all_text_features.norm(dim=-1, keepdim=True)
+
+        # Calculate cosine similarities and apply softmax
+        # Using temperature parameter to control the "sharpness" of the distribution
+        temperature = 0.1  # Lower values make the distribution more peaked
+        cosine_similarities = image_features @ all_text_features.T
+        all_text_probs = torch.softmax(cosine_similarities / temperature, dim=-1)
+
+    # Convert to dictionaries
+    all_results = {label: float(all_text_probs[0][i]) for i, label in enumerate(type_labels + scene_labels + expression_labels + clothing_labels)}
+
+    return all_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."
-)
+iface = gr.Blocks(title="NSFW Tagging with Finetuned CLIP")
+
+with iface:
+    gr.Markdown("# NSFW Tagging with Finetuned CLIP")
+    gr.Markdown("Upload an image to analyze its content across multiple NSFW categories.")
+    gr.Markdown("Uses [woweenie/open-clip-vit-h-nsfw-finetune](https://huggingface.co/woweenie/open-clip-vit-h-nsfw-finetune) finetuned on NSFW images.")
+    gr.Markdown("Disclaimer: This model is not perfect and may make mistakes. Use at your own risk.")
+
+    with gr.Tabs():
+        with gr.Tab("Combined Predictions"):
+            with gr.Row():
+                image_input1 = gr.Image(type="pil", label="Upload Image")
+                combined_output = gr.Label(label="Predicted Tags", num_top_classes=10)
+            predict_btn1 = gr.Button("Analyze")
+            predict_btn1.click(
+                fn=process_image_combined_tags,
+                inputs=image_input1,
+                outputs=combined_output
+            )
+
+        with gr.Tab("Categorical Predictions"):
+            with gr.Row():
+                image_input2 = gr.Image(type="pil", label="Upload Image")
+            with gr.Row():
+                type_output = gr.Label(label="Predicted Type", num_top_classes=3)
+                scene_output = gr.Label(label="Predicted Scene", num_top_classes=10) 
+                expression_output = gr.Label(label="Predicted Expression", num_top_classes=3)
+                clothing_output = gr.Label(label="Predicted Clothing", num_top_classes=5)
+            predict_btn2 = gr.Button("Analyze")
+            predict_btn2.click(
+                fn=process_image_separate_tags,
+                inputs=image_input2,
+                outputs=[type_output, scene_output, expression_output, clothing_output]
+            )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()