Updated requirements.txt and transformed pipeline to handler

README.md

@@ -36,8 +36,16 @@ HF_TOKEN = ""
 
 def predict(path_to_image: str = None):
     with open(path_to_image, "rb") as i:
-        b64 = base64.b64encode(i.read())
-    payload = {"inputs": {"image": b64.decode("utf-8")}}
+        image = i.read()
+    payload = {
+        "inputs": {"image": image},
+        "parameters": {
+                   "sample": True,
+                   "top_p":0.9,
+                   "min_length":5,
+                   "max_length":20
+        }
+    }
     response = r.post(
         ENDPOINT_URL, headers={"Authorization": f"Bearer {HF_TOKEN}"}, json=payload
     )

pipeline.py → handler.py

@@ -11,7 +11,7 @@ from torchvision.transforms.functional import InterpolationMode
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-class PreTrainedPipeline():
+class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
         self.model_path = os.path.join(path,'model_large_caption.pth') 
@@ -39,14 +39,14 @@ class PreTrainedPipeline():
             data (:obj:):
                 includes the input data and the parameters for the inference.
         Return:
-            A :obj:`dict`:. The object returned should be a dict of one list like [[{"label": 0.9939950108528137}]] containing :
+            A :obj:`dict`:. The object returned should be a dict of one list like {"caption": ["A hugging face at the office"]} containing :
                 - "caption": A string corresponding to the generated caption.
         """
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
 
-        # decode base64 image to PIL
-        image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
+       
+        image = Image.open(BytesIO(inputs['image']))
         image = self.transform(image).unsqueeze(0).to(device)   
         with torch.no_grad():
             caption = self.model.generate(

requirements.txt

@@ -1,4 +1,5 @@
-timm
-fairscale
+timm==0.4.12
+transformers==4.15.0
+fairscale==0.4.4
 requests
-Pillow
+Pillow