Create pipeline.py

pipeline.py

@@ -0,0 +1,58 @@
+from typing import Dict, List, Any
+import torch
+import torchvision
+import torch.nn as nn
+import torchvision.transforms as transforms
+
+from PIL import Image
+
+
+
+MODEL_PATH = './website_classifier.pth'
+
+# Function to load an image and perform the necessary transformations
+def process_image(image):
+    # Load Image
+    img = image.convert("RGB")
+    
+    # Apply transformations
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+    ])
+    
+    img_t = transform(img)
+    
+    # Convert to a batch of 1
+    img_u = torch.unsqueeze(img_t, 0)
+    
+    return img_u
+    
+class PreTrainedPipeline():
+    def __init__(self, path=""):
+        self.model = torchvision.models.resnet18(pretrained=True)
+        num_ftrs = self.model.fc.in_features
+        self.model.fc = nn.Linear(num_ftrs, 3)
+        self.transform = transforms.Compose(
+        [transforms.Resize((224, 224)),
+         transforms.ToTensor(),
+         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+        self.model.load_state_dict(torch.load(MODEL_PATH))
+        self.processor = process_image
+        self.classes = ['forum', 'general', 'marketplace']
+        self.classe_to_idx = {'forum': 0, 'general': 1, 'marketplace': 2}
+    
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        image = data.pop("inputs", data)
+
+        # process image
+        image = self.processor(image)
+
+        # run prediction
+        outputs = self.model.generate(image)
+        
+        # decode output
+        _, predicted = torch.max(outputs, 1)
+        prediction = self.classes[predicted[0]]
+        return {"class":prediction[0]}