Upload handler.py

handler.py

@@ -6,6 +6,7 @@ from typing import  Dict, List, Any
 import base64
 from io import BytesIO
 import os
+import boto3
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -25,6 +26,27 @@ class EndpointHandler():
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
         ])
 
+    def update_to_s3(self, image):
+        BUCKET_NAME = 'popwear-assets'
+        BUCKET_PREFIX_PATH = 'removebg'
+        ACCOUNT_ID = '18cc2282d0ee72171c1ea322ed22983c'
+        ACCESS_KEY_ID = '007f1852a377a2df43a21d5c8d54542e'
+        SECRET_ACCESS_KEY = 'db2658e2429950bb05e15afb6c53c8b7fd23ab9e1bf79cd42604c89f276068e4'
+        ENDPOINT_URL = f'https://{ACCOUNT_ID}.r2.cloudflarestorage.com'
+        bucket_postfix_path = f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
+        
+        image_url = f"https://assets.popwear.ai/{bucket_prefix_path}/{bucket_postfix_path}"
+        s3 = boto3.client(
+            's3',
+            endpoint_url=ENDPOINT_URL,
+            aws_access_key_id=ACCESS_KEY_ID,
+            aws_secret_access_key=SECRET_ACCESS_KEY,
+            region_name='auto'
+        )
+        image.seek(0)
+        s3.upload_fileobj(image, BUCKET_NAME, f"{BUCKET_PREFIX_PATH}/{bucket_postfix_path}")
+        return image_url
+
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         image = data.pop("inputs", data)
         # image = self.decode_base64_image(image)
@@ -35,7 +57,8 @@ class EndpointHandler():
         pred_pil = transforms.ToPILImage()(pred)
         mask = pred_pil.resize(image.size)
         image.putalpha(mask)
-        return image
+        image_url = self.update_to_s3(image)
+        return image_url
 
     def decode_base64_image(self, image_string):
         base64_image = base64.b64decode(image_string)