image moderation

app.py

@@ -1,13 +1,46 @@
+import os
+from PIL import Image
+import torch
+from torchvision import transforms
+from transformers import AutoProcessor, FocalNetForImageClassification
 import gradio as gr
 import numpy as np
 import random
 from diffusers import DiffusionPipeline
-import torch
 from huggingface_hub import InferenceClient
 import requests
-from PIL import Image
 from io import BytesIO
 
+# Paths and model setup
+image_folder = "path_to_your_image_folder"  # Specify the path to your image folder
+model_path = "MichalMlodawski/nsfw-image-detection-large"
+
+# List of jpg files in the folder
+jpg_files = [file for file in os.listdir(image_folder) if file.lower().endswith(".jpg")]
+
+if not jpg_files:
+    print("🚫 No jpg files found in folder:", image_folder)
+    exit()
+
+# Load the model and feature extractor
+feature_extractor = AutoProcessor.from_pretrained(model_path)
+model = FocalNetForImageClassification.from_pretrained(model_path)
+model.eval()
+
+# Image transformations
+transform = transforms.Compose([
+    transforms.Resize((512, 512)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Mapping from model labels to NSFW categories
+label_to_category = {
+    "LABEL_0": "Safe",
+    "LABEL_1": "Questionable",
+    "LABEL_2": "Unsafe"
+}
+
 # Device configuration
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -24,6 +57,9 @@ else:
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+# Initialize the InferenceClient
+client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+
 # Inference function for generating images
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
     if randomize_seed:
@@ -43,16 +79,6 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     
     return image
 
-# Examples for the text-to-image generation
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-# Initialize the InferenceClient
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
 # Respond function for the chatbot
 def respond(message, history, system_message, max_tokens, temperature, top_p):
     messages = [{"role": "system", "content": system_message}]
@@ -186,5 +212,22 @@ with gr.Blocks(css=css) as demo:
             image_moderation_result = gr.Textbox(label="Image Moderation Result")
             moderate_image_button.click(moderate_image, uploaded_image, image_moderation_result)
 
-if __name__ == "__main__": 
-    demo.launch()
+        with gr.TabItem("NSFW Classification"):
+            selected_image = gr.Image(type="pil", label="Upload Image for NSFW Classification")
+            classify_button = gr.Button("Classify Image")
+            classification_result = gr.Textbox(label="Classification Result")
+
+            def classify_nsfw(image):
+                image_tensor = transform(image).unsqueeze(0)
+                inputs = feature_extractor(images=image, return_tensors="pt")
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
+                    confidence, predicted = torch.max(probabilities, 1)
+                label = model.config.id2label[predicted.item()]
+                category = label_to_category.get(label, "Unknown")
+                return f"Label: {label}, Category: {category}, Confidence: {confidence.item() * 100:.2f}%"
+
+            classify_button.click(classify_nsfw, selected_image, classification_result)
+
+demo.launch()