Views
No views yet
microsoft/focalnet-base. My superpower is the lightning-fast classification of images into three categories:pip install transformers==4.37.2 torch==2.3.1 torchvision Pillow1import os
2from PIL import Image
3import torch
4from torchvision import transforms
5from transformers import AutoProcessor, FocalNetForImageClassification
6
7# Path to the folder with images
8image_folder = ""
9# Path to the model
10model_path = "MichalMlodawski/nsfw-image-detection-large"
11
12# List of jpg files in the folder
13jpg_files = [file for file in os.listdir(image_folder) if file.lower().endswith(".jpg")]
14
15# Check if there are jpg files in the folder
16if not jpg_files:
17 print("🚫 No jpg files found in folder:", image_folder)
18 exit()
19
20# Load the model and feature extractor
21feature_extractor = AutoProcessor.from_pretrained(model_path)
22model = FocalNetForImageClassification.from_pretrained(model_path)
23model.eval()
24
25# Image transformations
26transform = transforms.Compose([
27 transforms.Resize((512, 512)),
28 transforms.ToTensor(),
29 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
30])
31
32# Mapping from model labels to NSFW categories
33label_to_category = {
34 "LABEL_0": "Safe",
35 "LABEL_1": "Questionable",
36 "LABEL_2": "Unsafe"
37}
38
39# Processing and prediction for each image
40results = []
41for jpg_file in jpg_files:
42 selected_image = os.path.join(image_folder, jpg_file)
43 image = Image.open(selected_image).convert("RGB")
44 image_tensor = transform(image).unsqueeze(0)
45
46 # Process image using feature_extractor
47 inputs = feature_extractor(images=image, return_tensors="pt")
48
49 # Prediction using the model
50 with torch.no_grad():
51 outputs = model(**inputs)
52 probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
53 confidence, predicted = torch.max(probabilities, 1)
54
55 # Get the label from the model's configuration
56 label = model.config.id2label[predicted.item()]
57
58 results.append((jpg_file, label, confidence.item() * 100))
59
60# Display results
61print("🖼️ NSFW Classification Results 🖼️")
62print("=" * 40)
63for jpg_file, label, confidence in results:
64 category = label_to_category.get(label, "Unknown")
65 emoji = {"Safe": "✅", "Questionable": "⚠️", "Unsafe": "🔞"}.get(category, "❓")
66 confidence_bar = "🟩" * int(confidence // 10) + "⬜" * (10 - int(confidence // 10))
67
68 print(f"📄 File name: {jpg_file}")
69 print(f"🏷️ Model Label: {label}")
70 print(f"{emoji} NSFW Category: {category}")
71 print(f"🎯 Confidence: {confidence:.2f}% {confidence_bar}")
72 print(f"{'=' * 40}")
73
74print("🏁 Classification completed! 🎉")