Update handler.py

handler.py

@@ -24,12 +24,19 @@ class EndpointHandler():
         image=info.pop('image',data)
         image=base64.b64decode(image)
         raw_image=Image.open(BytesIO(image)).convert('RGB')
-        target=info.pop('target',data)
-        replacement=info.pop('replacement',data)
+        texts=info.pop('texts',data)
+        target=texts[0]
+        replacement=texts[1]
         if replacement=="":
-            return self.model.inference_remove(raw_image,target)
+            img=self.model.inference_remove(raw_image,target)
+            img_data=img.tobytes()
+            img_base64=base64.b64encode(img_data).decode('utf-8')
+            return {'image':img_base64}
         else:
-            return json.dumps({'image':self.model.inference_replace_sam(raw_image,target,replacement)})
+            img=self.model.inference_replace_sam(raw_image,target,replacement)
+            img_data=img.tobytes()
+            img_base64=base64.b64encode(img_data).decode('utf-8')
+            return {'image':img_base64}
             
     
 if __name__=="__main__":