Views
No views yet
1pip install transformers==4.37.2
2pip install torch==2.3.1
3pip install Pillow1import os
2from PIL import Image
3import torch
4from torchvision import transforms
5from transformers import AutoImageProcessor, MobileNetV2ForImageClassification
6
7# Path to the folder with images
8image_folder = ""
9# Path to the model
10model_path = "MichalMlodawski/open-closed-eye-classification-mobilev2"
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 image processor
21image_processor = AutoImageProcessor.from_pretrained(model_path)
22model = MobileNetV2ForImageClassification.from_pretrained(model_path)
23model.eval()
24
25# Image transformations
26transform = transforms.Compose([
27 transforms.Resize((256, 256)),
28 transforms.ToTensor()
29])
30
31# Processing and prediction for each image
32results = []
33for jpg_file in jpg_files:
34 selected_image = os.path.join(image_folder, jpg_file)
35 image = Image.open(selected_image).convert("RGB")
36 image_tensor = transform(image).unsqueeze(0)
37
38 # Process image using image_processor
39 inputs = image_processor(images=image, return_tensors="pt")
40
41 # Prediction using the model
42 with torch.no_grad():
43 outputs = model(**inputs)
44 probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
45 confidence, predicted = torch.max(probabilities, 1)
46
47 results.append((jpg_file, predicted.item(), confidence.item() * 100))
48
49# Display results
50print("🖼️ Image Classification Results 🖼️")
51print("=" * 40)
52
53for jpg_file, prediction, confidence in results:
54 emoji = "👁️" if prediction == 1 else "❌"
55 confidence_bar = "🟩" * int(confidence // 10) + "⬜" * (10 - int(confidence // 10))
56
57 print(f"📄 File name: {jpg_file}")
58 print(f"{emoji} Prediction: {'Open' if prediction == 1 else 'Closed'}")
59 print(f"🎯 Confidence: {confidence:.2f}% {confidence_bar}")
60 print(f"{'=' * 40}")
61
62print("🏁 Classification completed! 🎉")