new backend

backend.py

@@ -1,10 +1,11 @@
 import base64
+from io import BytesIO
 import numpy as np
 import cv2
 from segment_anything import sam_model_registry, SamPredictor
 import torch
 import time
-from flask import Flask, request
+from flask import Flask, Response, request
 
 # check torch version and if cuda is available
 print(torch.__version__)
@@ -24,7 +25,9 @@ app = Flask(__name__)
 def base64_to_image(base64_code):
     img_data = base64.b64decode(base64_code)
     img_array = np.fromstring(img_data, np.uint8)
-    img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)
+    print(img_array)
+    img = cv2.imdecode(img_array, -1)
+    print(img)
     return img
 
 
@@ -38,16 +41,19 @@ def index():
     if (content_type == 'application/json'):
         print('start calculate embedding')
         base64data = request.get_json()['imgurl']
-        print(base64data)
         image = base64_to_image(base64data)
         predictor.set_image(image)
         image_embedding = predictor.get_image_embedding().cpu().numpy()
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        print(time.time())
-        return image_embedding
+        byte_io = BytesIO()
+        np.save(byte_io, image_embedding)
+        byte_io.seek(0)
+        np.save("array.npy", image_embedding)
+        response = Response(byte_io, mimetype="application/octet-stream")
+        response.headers["Content-Length"] = 4194432
+        return response
     else:
         return 'Content-Type not supported!'
-
-
-# app.run(port=2020, host="0.0.0.0", debug=True)
-print('server starts working')
+    
+    
+print('server starts working')