""" ViT Seal Classifier - Evaluation and prediction module. """ import torch import torch.nn as nn from torchvision import transforms from transformers import ViTForImageClassification, ViTConfig from PIL import Image import os import json import numpy as np class ViTSealClassifier: def __init__(self, model_path='vit_seal_checker.pth'): self.model_path = model_path self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.model = None self.transform = None self.classes = ['Real', 'Fake'] self.is_loaded = False # Initialize transform self.transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) def load_model(self): """Load the trained ViT model.""" if self.is_loaded: return True # Try to download model if not present if not os.path.exists(self.model_path): try: from model_downloader import download_vit_model downloaded_path = download_vit_model() if downloaded_path: self.model_path = downloaded_path else: print(f"❌ Model file not found: {self.model_path}") print("Please run train_vit_seal_model.py first to train the model.") return False except ImportError: print(f"❌ Model file not found: {self.model_path}") print("Please run train_vit_seal_model.py first to train the model.") return False try: # Try to load pretrained ViT try: self.model = ViTForImageClassification.from_pretrained( 'google/vit-base-patch16-224', num_labels=2, ignore_mismatched_sizes=True ) except Exception as e: print(f"Using local ViT configuration: {e}") config = ViTConfig( image_size=224, patch_size=16, num_channels=3, num_labels=2, hidden_size=768, num_hidden_layers=12, num_attention_heads=12 ) self.model = ViTForImageClassification(config) # Load trained weights checkpoint = torch.load(self.model_path, map_location=self.device) self.model.load_state_dict(checkpoint['model_state_dict']) self.model = self.model.to(self.device) self.model.eval() self.is_loaded = True print(f"✅ ViT Seal Classifier loaded successfully!") # Load model info if available info_path = 'vit_model_info.json' if os.path.exists(info_path): with open(info_path, 'r') as f: model_info = json.load(f) print(f"Model accuracy: {model_info.get('best_accuracy', 'Unknown'):.2f}%") return True except Exception as e: print(f"❌ Error loading model: {e}") return False def predict_image(self, image_input): """ Predict if a seal/stamp is real or fake. Args: image_input: Can be PIL Image, image path (str), or numpy array Returns: dict: Prediction results with confidence scores """ if not self.load_model(): return { "error": "Model not loaded", "seal_status": "Error", "confidence": 0.0 } try: # Handle different input types if isinstance(image_input, str): # Image path if not os.path.exists(image_input): return { "error": f"Image file not found: {image_input}", "seal_status": "Error", "confidence": 0.0 } image = Image.open(image_input).convert('RGB') elif isinstance(image_input, Image.Image): # PIL Image image = image_input.convert('RGB') elif isinstance(image_input, np.ndarray): # Numpy array image = Image.fromarray(image_input).convert('RGB') else: return { "error": "Unsupported image input type", "seal_status": "Error", "confidence": 0.0 } # Preprocess image image_tensor = self.transform(image).unsqueeze(0).to(self.device) # Make prediction with torch.no_grad(): outputs = self.model(image_tensor) probabilities = torch.nn.functional.softmax(outputs.logits, dim=1) confidence, predicted = torch.max(probabilities, 1) # Get individual class probabilities real_prob = probabilities[0][0].item() fake_prob = probabilities[0][1].item() # Determine result is_real = predicted.item() == 0 seal_status = "Real" if is_real else "Fake" result = { "step": "Seal Verification", "status": "Pass" if is_real else "Fail", "reason": f"Seal classified as {seal_status.lower()} with {confidence.item():.2%} confidence", "seal_status": seal_status, "confidence": confidence.item(), "probabilities": { "Real": real_prob, "Fake": fake_prob }, "prediction_text": f"{seal_status} ✅" if is_real else f"{seal_status} ❌" } return result except Exception as e: return { "error": f"Prediction error: {str(e)}", "seal_status": "Error", "confidence": 0.0 } def predict_multiple_seals(self, seal_images): """ Predict multiple seals and return combined result. Args: seal_images: List of image inputs Returns: dict: Combined prediction results """ if not seal_images: return { "error": "No seal images provided", "seal_status": "Error", "confidence": 0.0 } predictions = [] for i, seal_image in enumerate(seal_images): result = self.predict_image(seal_image) result['seal_index'] = i + 1 predictions.append(result) # Improved combine results - more lenient approach overall_real_count = sum(1 for p in predictions if p.get('seal_status') == 'Real') overall_fake_count = len(predictions) - overall_real_count # Calculate average confidence valid_predictions = [p for p in predictions if 'confidence' in p and p['confidence'] > 0] avg_confidence = np.mean([p['confidence'] for p in valid_predictions]) if valid_predictions else 0.0 # More lenient decision logic: Pass if majority are real OR if confidence is reasonable total_seals = len(predictions) real_ratio = overall_real_count / total_seals if total_seals > 0 else 0 # Pass if: # 1. Majority (>50%) are real, OR # 2. At least one real seal with high confidence (>80%), OR # 3. Average confidence is reasonable (>60%) regardless of classification high_confidence_real = any(p.get('seal_status') == 'Real' and p.get('confidence', 0) > 0.8 for p in predictions) is_overall_real = (real_ratio > 0.5) or high_confidence_real or (avg_confidence > 0.6) combined_result = { "step": "Seal Verification", "status": "Pass" if is_overall_real else "Fail", "reason": f"Found {overall_real_count} real seals, {overall_fake_count} fake seals", "seal_status": "Real" if is_overall_real else "Fake", "confidence": avg_confidence, "seal_count": len(predictions), "real_count": overall_real_count, "fake_count": overall_fake_count, "individual_predictions": predictions } return combined_result def create_dummy_prediction(self, confidence=0.85): """Create a dummy prediction for demo purposes.""" import random is_real = random.choice([True, False]) seal_status = "Real" if is_real else "Fake" return { "step": "Seal Verification", "status": "Pass" if is_real else "Fail", "reason": f"Seal classified as {seal_status.lower()} (demo mode)", "seal_status": seal_status, "confidence": confidence, "probabilities": { "Real": confidence if is_real else 1 - confidence, "Fake": 1 - confidence if is_real else confidence }, "prediction_text": f"{seal_status} ✅" if is_real else f"{seal_status} ❌", "demo_mode": True } # Utility functions def evaluate_seal_folder(folder_path, model_path='vit_seal_checker.pth'): """Evaluate all images in a folder.""" classifier = ViTSealClassifier(model_path) if not os.path.exists(folder_path): print(f"Folder not found: {folder_path}") return image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] if not image_files: print(f"No image files found in {folder_path}") return print(f"Evaluating {len(image_files)} images in {folder_path}...") results = [] for image_file in image_files: image_path = os.path.join(folder_path, image_file) result = classifier.predict_image(image_path) result['filename'] = image_file results.append(result) status_icon = "✅" if result.get('seal_status') == 'Real' else "❌" confidence = result.get('confidence', 0) print(f"{image_file}: {result.get('seal_status', 'Error')} {status_icon} ({confidence:.2%})") return results def test_single_image(image_path, model_path='vit_seal_checker.pth'): """Test classification on a single image.""" classifier = ViTSealClassifier(model_path) result = classifier.predict_image(image_path) print(f"Image: {image_path}") print(f"Prediction: {result}") return result if __name__ == "__main__": # Test the classifier if model exists if os.path.exists('vit_seal_checker.pth'): print("🧪 Testing ViT Seal Classifier...") # Test with seal dataset if available test_folders = ['seal_dataset/val/real', 'seal_dataset/val/fake'] for folder in test_folders: if os.path.exists(folder): print(f"\n📁 Testing folder: {folder}") evaluate_seal_folder(folder) else: print("❌ Model not found. Please run train_vit_seal_model.py first.") # Create dummy prediction for testing classifier = ViTSealClassifier() dummy_result = classifier.create_dummy_prediction() print(f"Demo prediction: {dummy_result}")