Upload 6 files

custom_classifier/__init__.py

@@ -0,0 +1,2 @@
+from .configuration import CustomModelConfig
+from .model import CustomClassifier

custom_classifier/configuration.py

@@ -0,0 +1,11 @@
+from transformers import PretrainedConfig
+
+class CustomModelConfig(PretrainedConfig):
+    model_type = "custom"
+    input_size = (3, 128, 128)
+    num_classes = 2
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.input_size = kwargs.get("input_size", self.input_size)
+        self.num_classes = kwargs.get("num_classes", self.num_classes)

custom_classifier/model.py

@@ -0,0 +1,87 @@
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from .configuration import CustomModelConfig
+from torchvision import transforms
+from PIL import Image
+import sys
+
+class CustomModel(nn.Module):
+    def __init__(self, input_shape, num_classes):
+        super(CustomModel, self).__init__()
+
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(in_channels=input_shape[0], out_channels=32, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(32),
+            nn.MaxPool2d(kernel_size=2),
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(64),
+            nn.MaxPool2d(kernel_size=2),
+            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(128),
+            nn.MaxPool2d(kernel_size=2),
+            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(128),
+            nn.MaxPool2d(kernel_size=2)
+        )
+
+        self.fc_layers = nn.Sequential(
+            nn.Flatten(),
+            nn.Dropout(0.5),
+            nn.Linear(128 * (input_shape[1] // 16) * (input_shape[2] // 16), 512),
+            nn.ReLU(),
+            nn.BatchNorm1d(512),
+            nn.Dropout(0.5),
+            nn.Linear(512, num_classes)
+        )
+
+    def forward(self, x):
+        x = self.conv_layers(x)
+        x = self.fc_layers(x)
+        return x
+
+class CustomClassifier(PreTrainedModel):
+    config_class = CustomModelConfig
+
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = CustomModel(config.input_size, config.num_classes)
+        self.preprocess = transforms.Compose([
+                transforms.Resize((128, 128)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+        self.classes = ['cat', 'dog']
+
+    def forward(self, x):
+        try:
+            x = Image.open(x).convert("RGB")
+        except Exception as e:
+            raise Exception(f"Error: Unable to load image file {x}. Check if the file exists or is in the right format. Details: {e}")
+
+        x = self.preprocess(x).unsqueeze(0)
+        
+
+        return self.model(x)
+    
+    def predict(self, x, get_class=False):
+        self.eval()
+        with torch.no_grad():
+            outputs = self.forward(x)
+            probabilities = torch.nn.functional.softmax(outputs, dim=1)
+            if not get_class:
+                return {
+                    "cat": round(probabilities[0][0].item(), 3),
+                    "dog": round(probabilities[0][1].item(), 3)
+                }
+            else:
+                
+                return self.classes[probabilities.argmax(dim=1).item()]
+        
+
+