Update api_server.py

api_server.py

@@ -35,7 +35,7 @@ if load_type == 'local':
         
     model = YOLO(model_path)
 
-    print("===============LOAD YOLO MODEL DONE =============")
+    print("***** 1. LOAD YOLO MODEL DONE *****")
     #model.eval()  # 設定模型為推理模式
 elif load_type == 'remote_hub_download':
     from huggingface_hub import hf_hub_download
@@ -86,10 +86,10 @@ 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()
@@ -121,10 +121,10 @@ def predict():
     except Exception as e:
         return jsonify({'error': str(e)}), 400
         
-    print("***** Start YOLO predict *****")
+    print("***** 2. Start YOLO predict *****")
     # Make a prediction using YOLO
     results = model(image_data)
-    print ("***** YOLO predict result:",results,"********")
+    print ("===== YOLO predict result:",results,"=====")
     print("***** YOLO predict DONE *****")
     
     check_memory_usage()
@@ -149,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 label_names: {label_names}======")
+        print(f"====== 3. YOLO label_names: {label_names}======")
 
         element_counts = Counter(label_names)
 
@@ -158,7 +158,7 @@ def predict():
             yolo_path = f"{YOLO_DIR}/{message_id}/{element}"
             yolo_file = get_jpg_files(yolo_path)
 
-            print(f"======YOLO result:{yolo_path}======")
+            print(f"***** 處理:{yolo_path} *****")
 
             if len(yolo_file) == 0:
                 print(f"警告：{element} 沒有找到相關的 JPG 檔案")
@@ -167,10 +167,10 @@ def predict():
             element_list.append(element)
             
             for yolo_img in yolo_file: # 每張切圖yolo_img
-                print("***** START CLIP *****")
+                print("***** 4. 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"===== CLIP result:{top_k_words} =====\n")
 
             # if element_counts[element] > 1: #某隻角色的數量>1
             #     yolo_path = f"{YOLO_DIR}/{message_id}/{element}"
@@ -201,7 +201,7 @@ def predict():
         #     'description': element_list
         # }
         
-        return top_k_words#jsonify(response_data)
+        return jsonify(top_k_words), 200 #jsonify(response_data)
             
         # for label_name in label_names:
         #     yolo_file=f"{YOLO_DIR}/{message_id}/{label_name}/im.jpg.jpg"