Upload model.py with huggingface_hub

model.py

@@ -1,6 +1,10 @@
-from transformers import PreTrainedModel, PretrainedConfig
+from transformers import PreTrainedModel, PretrainedConfig, ImageProcessingMixin
+import torch
 import torch.nn as nn
 from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
+from torchvision import transforms
+from PIL import Image
+import numpy as np
 
 class CheckboxConfig(PretrainedConfig):
     model_type = "checkbox-classifier"
@@ -9,6 +13,44 @@ class CheckboxConfig(PretrainedConfig):
         super().__init__(num_labels=num_labels, **kwargs)
         self.dropout_rate = dropout_rate
 
+class CheckboxImageProcessor(ImageProcessingMixin):
+    """Simple image processor for checkbox classifier"""
+    
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.size = {"height": 128, "width": 128}
+        self.image_mean = [0.485, 0.456, 0.406]
+        self.image_std = [0.229, 0.224, 0.225]
+        
+        self.transform = transforms.Compose([
+            transforms.Resize((self.size["height"], self.size["width"])),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=self.image_mean, std=self.image_std)
+        ])
+    
+    def preprocess(self, images, **kwargs):
+        """Preprocess images for model input"""
+        if not isinstance(images, list):
+            images = [images]
+        
+        processed = []
+        for image in images:
+            if isinstance(image, str):
+                image = Image.open(image).convert('RGB')
+            elif isinstance(image, np.ndarray):
+                image = Image.fromarray(image).convert('RGB')
+            elif not isinstance(image, Image.Image):
+                raise ValueError(f"Unsupported image type: {type(image)}")
+            
+            processed.append(self.transform(image))
+        
+        # Stack into batch
+        pixel_values = torch.stack(processed)
+        return {"pixel_values": pixel_values}
+    
+    def __call__(self, images, **kwargs):
+        return self.preprocess(images, **kwargs)
+
 class CheckboxClassifier(PreTrainedModel):
     config_class = CheckboxConfig
     
@@ -16,7 +58,7 @@ class CheckboxClassifier(PreTrainedModel):
         super().__init__(config)
         self.num_labels = config.num_labels
         
-        self.backbone = efficientnet_v2_s(weights=EfficientNet_V2_S_Weights.IMAGENET1K_V1)
+        self.backbone = efficientnet_v2_s(weights=None)  # Don't load pretrained weights here
         num_features = self.backbone.classifier[1].in_features
         
         self.backbone.classifier = nn.Sequential(
@@ -32,6 +74,15 @@ class CheckboxClassifier(PreTrainedModel):
             nn.Linear(256, config.num_labels)
         )
     
-    def forward(self, pixel_values):
+    def forward(self, pixel_values, labels=None):
         outputs = self.backbone(pixel_values)
-        return {"logits": outputs}
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(outputs, labels)
+            
+        return {
+            "loss": loss,
+            "logits": outputs,
+        }