Upload MnistPipe

MyPipe.py

@@ -0,0 +1,65 @@
+
+from transformers import Pipeline
+import requests
+from PIL import Image
+import torchvision.transforms as transforms
+import torch
+
+class MnistPipe(Pipeline):
+    def __init__(self,**kwargs):
+
+      # self.tokenizer = (...) # code if you want to instantiate more parameters
+
+      Pipeline.__init__(self,**kwargs) # self.model automatically instantiated here
+
+      self.transform = transforms.Compose(
+                              [transforms.ToTensor(),
+                              transforms.Resize((28,28), antialias=True)
+                              ])
+
+    def _sanitize_parameters(self, **kwargs):
+        # will make sure where each parameter goes
+        preprocess_kwargs = {}
+        postprocess_kwargs = {}
+        if "download" in kwargs:
+            preprocess_kwargs["download"] = kwargs["download"]
+        if "clean_output" in kwargs :
+          postprocess_kwargs["clean_output"] = kwargs["clean_output"]
+        return preprocess_kwargs, {}, postprocess_kwargs
+
+    def preprocess(self, inputs, download=False):
+        if download == True :
+          # call download_img method and name image as "image.png"
+          self.download_img(inputs)
+          inputs = "image.png"
+
+        # we open and process the image
+        img = Image.open(inputs)
+        gray = img.convert('L')
+        tensor = self.transform(gray)
+        tensor = tensor.unsqueeze(0)
+        return tensor
+
+    def _forward(self, tensor):
+        with torch.no_grad():
+            # the model has been automatically instantiated
+            # in the __init__ method
+            out = self.model(tensor)
+        return out
+
+    def postprocess(self, out, clean_output=True):
+        if clean_output ==True :
+          label = torch.argmax(out,axis=-1) # get class
+          label = label.tolist()[0]
+          return label
+        else :
+          return out
+
+    def download_img(self,url):
+      # if download = True download image and name it image.png
+      response = requests.get(url, stream=True)
+
+      with open("image.png", "wb") as f:
+          for chunk in response.iter_content(chunk_size=8192):
+              f.write(chunk)
+      print("image saved as image.png")

config.json

@@ -1,13 +1,24 @@
 {
+  "_name_or_path": "not-lain/mycustomrepo",
   "architectures": [
     "MnistModel"
   ],
   "auto_map": {
-    "AutoConfig": "MyConfig.MnistConfig",
-    "AutoModelForImageClassification": "MyModel.MnistModel"
+    "AutoConfig": "not-lain/mycustomrepo--MyConfig.MnistConfig",
+    "AutoModelForImageClassification": "not-lain/mycustomrepo--MyModel.MnistModel"
   },
   "conv1": 10,
   "conv2": 20,
+  "custom_pipelines": {
+    "image-classification": {
+      "impl": "MyPipe.MnistPipe",
+      "pt": [
+        "AutoModelForImageClassification"
+      ],
+      "tf": [],
+      "type": "image"
+    }
+  },
   "model_type": "MobileNetV1",
   "torch_dtype": "float32",
   "transformers_version": "4.39.0.dev0"