Update api_server.py

api_server.py

@@ -8,14 +8,14 @@ from ultralytics import YOLO
 import io
 import base64
 import uuid
-
-# Disable tensorflow warnings
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
-
+import glob
 from tensorflow import keras
 from flask import Flask, jsonify, request, render_template, send_file
 import torch
 
+# Disable tensorflow warnings
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+
 load_type = 'local'
 
 MODEL_NAME = "yolo11_detect_best_241018_1.pt"
@@ -50,12 +50,29 @@ else:
     raise AssertionError('No load type is specified!')
 
 
+# image to base64
 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
 
 
+# 抓取指定路徑下的所有 JPG 檔案
+def get_jpg_files(path):
+  """
+  Args:
+    path: 要搜尋的目錄路徑。
+
+  Returns:
+    一個包含所有 JPG 檔案路徑的列表。
+  """
+  return glob.glob(os.path.join(path, "*.jpg"))
+
+# 使用範例
+image_folder = '/content/drive/MyDrive/chiikawa'  # 替換成你的目錄路徑
+jpg_files = get_jpg_files(image_folder)
+
+
 # Initialize the Flask application
 app = Flask(__name__)
 
@@ -65,18 +82,13 @@ app = Flask(__name__)
 def predict():
 
     user_id = request.args.get('user_id') 
-    
-    # 生成一個唯一的 message_id
-    message_id = str(uuid.uuid4())
+    file = request.files['image']
+    message_id = str(uuid.uuid4()) # 生成一個唯一的 message_id
     
     if 'image' not in request.files:
         # Handle if no file is selected
         return 'No file selected'
 
-    start_time = time.time()
-
-    file = request.files['image']
-
     # 讀取圖像
     try:
         image_data = Image.open(file)
@@ -90,20 +102,6 @@ def predict():
     if results is None or len(results) == 0:
         return jsonify({'error': 'No results from YOLO model'}), 400
 
-    # # 渲染推理結果到圖像
-    # img_with_boxes = results[0].plot()  # 使用 results[0]，假設只有一張圖像作推理
-
-    # # 將 numpy array 轉換為 PIL Image
-    # img = Image.fromarray(img_with_boxes)
-
-    # # 儲存圖片到內存緩衝區
-    # img_io = io.BytesIO()
-    # img.save(img_io, 'PNG')
-    # img_io.seek(0)
-
-    # # 返回處理後的圖像
-    # return send_file(img_io, mimetype='image/png')
-
     saved_images = []
     
     # 儲存辨識後的圖片到指定資料夾
@@ -115,22 +113,55 @@ def predict():
         labels = result.boxes.cls  # Get predicted label IDs
         label_names = [model.names[int(label)] for label in labels]  # Convert to names
 
-        encoded_images=[]
+        element_counts = Counter(my_list)
 
-        for label_name in label_names:
-            output_file=f"{YOLO_DIR}/{message_id}/{label_name}/im.jpg.jpg"
-            # 將圖片轉換為 base64 編碼
-            encoded_images.append(image_to_base64(output_file))
-            
+        encoded_images=[]
+        element_list =[]
+
+        for element, count in element_counts.items():
+            if element_counts[element] > 1: #某隻角色的數量>1
+                output_path = f"{YOLO_DIR}/{message_id}/{element}"
+                output_file = get_jpg_files(output_path)
+                element_list.append(element)
+
+                for output_img in output_file: # 取得每張圖的路徑
+                    encoded_images.append(image_to_base64(output_img))
+
+            else  : #某隻角色的數量=1
+                output_path = f"{YOLO_DIR}/{message_id}/{element}/im.jpg.jpg"
+                encoded_images.append(image_to_base64(output_path))
+                element_list.append(element)
+                
         # 建立回應資料
         response_data = {
             'message_id': message_id,
             'images': encoded_images,
-            'description': label_names
+            'description': element_list
         }
-    return jsonify(response_data)
+        
+        return jsonify(response_data)
+        
+        # for label_name in label_names:
+        #     output_file=f"{YOLO_DIR}/{message_id}/{label_name}/im.jpg.jpg"
+        #     # 將圖片轉換為 base64 編碼
+        #     encoded_images.append(image_to_base64(output_file))
+
+        
+    # # 渲染推理結果到圖像
+    # img_with_boxes = results[0].plot()  # 使用 results[0]，假設只有一張圖像作推理
+
+    # # 將 numpy array 轉換為 PIL Image
+    # img = Image.fromarray(img_with_boxes)
+
+    # # 儲存圖片到內存緩衝區
+    # img_io = io.BytesIO()
+    # img.save(img_io, 'PNG')
+    # img_io.seek(0)
 
+    # # 返回處理後的圖像
+    # return send_file(img_io, mimetype='image/png')
 
+    
     # # dictionary is not a JSON: https://www.quora.com/What-is-the-difference-between-JSON-and-a-dictionary
     # # flask.jsonify vs json.dumps https://sentry.io/answers/difference-between-json-dumps-and-flask-jsonify/
     # # The flask.jsonify() function returns a Response object with Serializable JSON and content_type=application/json.