Update api_server.py

api_server.py

@@ -14,6 +14,7 @@ from flask import Flask, jsonify, request, render_template, send_file
 import torch
 from collections import Counter
 from clip_model import ClipModel
+import psutil
 
 # Disable tensorflow warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
@@ -75,7 +76,6 @@ def get_jpg_files(path):
 # image_folder = '/content/drive/MyDrive/chiikawa'  # 替換成你的目錄路徑
 # jpg_files = get_jpg_files(image_folder)
 
-import psutil
 
 def check_memory_usage():
     # Get memory details
@@ -86,20 +86,21 @@ def check_memory_usage():
     used_memory = memory_info.used / (1024 * 1024)
     memory_usage_percent = memory_info.percent
     
-    print(f"Total Memory: {total_memory:.2f} MB")
-    print(f"Available Memory: {available_memory:.2f} MB")
-    print(f"Used Memory: {used_memory:.2f} MB")
-    print(f"Memory Usage (%): {memory_usage_percent}%")
+    print(f"^^^^^^Total Memory: {total_memory:.2f} MB^^^^^^")
+    print(f"^^^^^^Available Memory: {available_memory:.2f} MB^^^^^^")
+    print(f"^^^^^^Used Memory: {used_memory:.2f} MB^^^^^^")
+    print(f"^^^^^^Memory Usage (%): {memory_usage_percent}%^^^^^^")
 
 # Run the function
 check_memory_usage()
-    
+
+
 # Initialize the Flask application
 app = Flask(__name__)
 # Initialize the ClipModel at the start
 clip_model = ClipModel()
 
-check_memory_usage()
+
 
 # API route for prediction(YOLO)
 @app.route('/predict', methods=['POST'])
@@ -125,6 +126,7 @@ def predict():
     results = model(image_data)
     print ("***** YOLO predict result:",results,"********")
     print("***** YOLO predict DONE *****")
+    
     check_memory_usage()
     
     # 檢查 YOLO 是否返回了有效的結果
@@ -147,7 +149,7 @@ def predict():
         labels = result.boxes.cls  # Get predicted label IDs
         label_names = [model.names[int(label)] for label in labels]  # Convert to names
 
-        print(f"======YOLO result: {label_names}======")
+        print(f"======YOLO label_names: {label_names}======")
 
         element_counts = Counter(label_names)
 
@@ -165,10 +167,10 @@ def predict():
             element_list.append(element)
             
             for yolo_img in yolo_file: # 每張切圖yolo_img
-                print("*****START CLIP *****")
+                print("***** START CLIP *****")
                 top_k_words.append(clip_model.clip_result(yolo_img)) # CLIP預測3個結果(top_k_words)
                 #encoded_images.append(image_to_base64(yolo_img))
-                print(f"**{yolo_img}:{top_k_words}**\n")
+                print(f"{yolo_img}:{top_k_words}\n")
 
             # if element_counts[element] > 1: #某隻角色的數量>1
             #     yolo_path = f"{YOLO_DIR}/{message_id}/{element}"
@@ -181,18 +183,18 @@ def predict():
             #     yolo_path = f"{YOLO_DIR}/{message_id}/{element}/im.jpg.jpg"
             #     encoded_images.append(image_to_base64(yolo_path))
                 
-        # 建立回應資料
-        response_data = {
-            'message_id': message_id,
-            'description': element_list,
-            'images': [
-                {
-                    'encoded_image': encoded_image,
-                    'description_list': top_k_words
-                }
-                for encoded_image, description_list in zip(encoded_images, top_k_words)
-            ]
-        }
+        ## 建立回應資料
+        # response_data = {
+        #     'message_id': message_id,
+        #     'description': element_list,
+        #     'images': [
+        #         {
+        #             'encoded_image': encoded_image,
+        #             'description_list': top_k_words
+        #         }
+        #         for encoded_image, description_list in zip(encoded_images, top_k_words)
+        #     ]
+        # }
         # response_data = {
         #     'message_id': message_id,
         #     'images': encoded_images,