Update handler.py

handler.py

@@ -1,131 +1,50 @@
-import os
-import torch
-import torch.nn as nn
-import torchvision.models as models
-import torchvision.transforms as transforms
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 from PIL import Image
-import json
-import base64
-from io import BytesIO
-from typing import Dict, List, Any, Union
-
-# Define the model architecture based on EfficientNetV2-S
-class AIDetectorModel(nn.Module):
-    def __init__(self):
-        super(AIDetectorModel, self).__init__()
-        # Load EfficientNetV2-S as base model
-        self.base_model = models.efficientnet_v2_s(weights=None)
-        
-        # Replace classifier with custom layers
-        self.base_model.classifier = nn.Sequential(
-            nn.Linear(self.base_model.classifier[1].in_features, 1024),
-            nn.ReLU(),
-            nn.Dropout(p=0.3),
-            nn.Linear(1024, 512),
-            nn.ReLU(),
-            nn.Dropout(p=0.3),
-            nn.Linear(512, 2)  # 2 classes: real or AI-generated
-        )
-    
-    def forward(self, x):
-        return self.base_model(x)
+import torch
 
-# This is the handler that will be used by the Hugging Face Inference API
-class Handler:
-    def __init__(self):
-        self.initialized = False
-        self.model = None
-        self.device = None
-        self.transform = None
-    
-    def initialize(self, context=None):
-        """Initialize the handler with model and preprocessing"""
+class EndpointHandler:
+    def __init__(self, model_dir):
+        self.model = AutoModelForImageClassification.from_pretrained(model_dir)
+        self.processor = AutoFeatureExtractor.from_pretrained(model_dir)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"Using device: {self.device}")
-        
-        # Initialize the model
-        self.model = AIDetectorModel()
+        self.model.to(self.device)
         
-        # Load the trained weights
-        model_path = os.path.join(os.getcwd(), "best_model_improved.pth")
+    def __call__(self, data):
+        """
+        Args:
+            data: Image data in binary format
+        Returns:
+            Prediction result as a dictionary
+        """
         try:
-            # Try to load with strict=True first
-            self.model.load_state_dict(torch.load(model_path, map_location=self.device))
-            print(f"Model loaded successfully from {model_path}")
-        except Exception as e:
-            print(f"Error with strict loading: {e}")
-            print("Trying with strict=False...")
-            # If that fails, try with strict=False
-            self.model.load_state_dict(torch.load(model_path, map_location=self.device), strict=False)
-            print("Model loaded with strict=False")
+            # Load and process image
+            image = Image.open(data).convert("RGB")
+            inputs = self.processor(images=image, return_tensors="pt")
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
             
-        self.model.to(self.device)
-        self.model.eval()  # Set to evaluation mode
-        
-        # Define image transformations - same as used in training
-        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])
-        ])
-        
-        self.initialized = True
-        print("Model initialization complete")
-        return self
-    
-    def preprocess(self, data: Union[Dict, List, str, bytes]) -> torch.Tensor:
-        """Process input data for model inference"""
-        images = []
-        
-        # Handle different input formats
-        if isinstance(data, dict):
-            # API format where data is a dictionary
-            if "inputs" in data:
-                data = data["inputs"]
-            elif "image" in data:
-                data = data["image"]
-        
-        # Convert to list for batch processing
-        if not isinstance(data, list):
-            data = [data]
-        
-        for item in data:
-            # Process each item based on type
-            if isinstance(item, str):
-                if os.path.isfile(item):
-                    # Local file path
-                    image = Image.open(item).convert('RGB')
-                elif item.startswith("http"):
-                    # URL
-                    from urllib.request import urlopen
-                    image = Image.open(BytesIO(urlopen(item).read())).convert('RGB')
-                elif item.startswith("data:image"):
-                    # Base64 image with header
-                    image_data = item.split(",")[1]
-                    image = Image.open(BytesIO(base64.b64decode(image_data))).convert('RGB')
-                else:
-                    # Assume base64 encoded image
-                    try:
-                        image = Image.open(BytesIO(base64.b64decode(item))).convert('RGB')
-                    except Exception:
-                        # If not base64, try as file path again
-                        image = Image.open(item).convert('RGB')
-            elif isinstance(item, bytes):
-                # Raw bytes
-                image = Image.open(BytesIO(item)).convert('RGB')
-            elif isinstance(item, Image.Image):
-                # Already a PIL Image
-                image = item.convert('RGB')
-            else:
-                raise ValueError(f"Unsupported input type: {type(item)}")
+            # Get prediction
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                logits = outputs.logits
+                probabilities = torch.nn.functional.softmax(logits, dim=-1)
+                
+            # Get prediction class and confidence
+            predicted_class_idx = probabilities.argmax().item()
+            confidence = probabilities[0][predicted_class_idx].item()
             
-            # Apply transformations
-            image_tensor = self.transform(image)
-            images.append(image_tensor)
-        
-        # Stack tensors for batch processing
-        batch = torch.stack(images).to(self.device)
-        return batch
-    
-    def inference(self, input_batch: torch.Tensor) -> List[Dict[str, Any]]:
-        """Run model inference on the input batch"""
+            # Get class labels
+            id2label = self.model.config.id2label
+            predicted_class = id2label[predicted_class_idx]
+            
+            # Return results
+            return {
+                "predicted_class": predicted_class,
+                "confidence": confidence,
+                "all_probabilities": {
+                    id2label[i]: prob.item() 
+                    for i, prob in enumerate(probabilities[0])
+                }
+            }
+            
+        except Exception as e:
+            return {"error": str(e)}