Update api_server.py

api_server.py

@@ -6,6 +6,7 @@ import torchvision.transforms as transforms
 from pathlib import Path
 from ultralytics import YOLO
 import io
+import base64
 
 # Disable tensorflow warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
@@ -47,6 +48,13 @@ elif load_type == 'remote_hub_from_pretrained':
 else:
     raise AssertionError('No load type is specified!')
 
+
+def image_to_base64(image_path):
+    with open(image_path, "rb") as image_file:
+        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
+        return encoded_string
+
+
 # Initialize the Flask application
 app = Flask(__name__)
 
@@ -54,6 +62,9 @@ app = Flask(__name__)
 # API route for prediction(YOLO)
 @app.route('/predict', methods=['POST'])
 def predict():
+
+    user_id = request.args.get('user_id') 
+    
     if 'image' not in request.files:
         # Handle if no file is selected
         return 'No file selected'
@@ -93,28 +104,26 @@ def predict():
     
     # 儲存辨識後的圖片到指定資料夾
     for result in results:
-        # 保存完整圖片
-        yolo_path = f'{YOLO_DIR}/results_{Path(result.path).name}'
-        result.save(yolo_path)
-        saved_images.append(yolo_path)
-
-        # 保存裁剪後的圖片（僅當 save_crop 返回有效的路徑時才加入）
-        cropped_images = result.save_crop(YOLO_DIR)  # 有些 YOLO 版本 save_crop 不返回值
-        if cropped_images:  # 確保不會對 None 進行迭代
-            if isinstance(cropped_images, list):  # 如果它返回一個列表
-                saved_images.extend(cropped_images)
-            else:
-                saved_images.append(cropped_images)
-
-    end_time = time.time()
-    inference_time = end_time - start_time
-
-    # 返回辨識結果的文件路徑以及推理時間
-    return jsonify({
-        'saved_images': saved_images,
-        'inference_time': inference_time
-    }), 200
-
+        # 保存圖片
+        result.save_crop(f"{YOLO_DIR}/{user_id}")
+        
+        num_detections = len(result.boxes)  # Get the number of detections
+        labels = result.boxes.cls  # Get predicted label IDs
+        label_names = [model.names[int(label)] for label in labels]  # Convert to names
+
+        encoded_images=[]
+
+        for label_name in label_names:
+            output_file=f"{YOLO_DIR}/{user_id}/{label_name}/im.jpg.jpg"
+            # 將圖片轉換為 base64 編碼
+            encoded_images.append(image_to_base64(output_file))
+            
+        # 建立回應資料
+        response_data = {
+            'images': encoded_images,
+            'description': label_names
+        }
+    return jsonify(response_data)
 
 
     # # dictionary is not a JSON: https://www.quora.com/What-is-the-difference-between-JSON-and-a-dictionary