Update api_server.py

api_server.py

@@ -13,13 +13,14 @@ from tensorflow import keras
 from flask import Flask, jsonify, request, render_template, send_file
 import torch
 from collections import Counter
+from clip import clip_result
 
 # Disable tensorflow warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 
 load_type = 'local'
 
-MODEL_NAME = "yolo11_detect_best_241024_1.pt"
+MODEL_YOLO = "yolo11_detect_best_241024_1.pt"
 MODEL_DIR = "./artifacts/models"
 YOLO_DIR = "./artifacts/yolo"
 #REPO_ID = "1vash/mnist_demo_model"
@@ -27,7 +28,7 @@ YOLO_DIR = "./artifacts/yolo"
 # Load the saved YOLO model into memory
 if load_type == 'local':
     # 本地模型路徑
-    model_path = f'{MODEL_DIR}/{MODEL_NAME}'
+    model_path = f'{MODEL_DIR}/{MODEL_YOLO}'
     if not os.path.exists(model_path):
         raise FileNotFoundError(f"Model file not found at {model_path}")
         
@@ -37,7 +38,7 @@ elif load_type == 'remote_hub_download':
     from huggingface_hub import hf_hub_download
 
     # 從 Hugging Face Hub 下載模型
-    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_NAME)
+    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_YOLO)
     model = torch.load(model_path)
     #model.eval()
 elif load_type == 'remote_hub_from_pretrained':
@@ -45,7 +46,7 @@ elif load_type == 'remote_hub_from_pretrained':
     os.environ['TRANSFORMERS_CACHE'] = str(Path(MODEL_DIR).absolute())
     from huggingface_hub import from_pretrained
 
-    model = from_pretrained(REPO_ID, filename=MODEL_NAME, cache_dir=MODEL_DIR)
+    model = from_pretrained(REPO_ID, filename=MODEL_YOLO, cache_dir=MODEL_DIR)
     #model.eval()
 else:
     raise AssertionError('No load type is specified!')
@@ -118,56 +119,54 @@ def predict():
 
         encoded_images=[]
         element_list =[]
+        top_k_words =[]
 
         for element, count in element_counts.items():
             
-            output_path = f"{YOLO_DIR}/{message_id}/{element}"
-            output_file = get_jpg_files(output_path)
+            yolo_path = f"{YOLO_DIR}/{message_id}/{element}"
+            yolo_file = get_jpg_files(yolo_path)
 
             element_list.append(element)
             
-            for output_img in output_file: # 取得每張圖的路徑
-                encoded_images.append(image_to_base64(output_img))
+            for yolo_img in yolo_file: # 每張切圖yolo_img
+                top_k_words.append(clip_result(yolo_img)) # CLIP預測3個結果(top_k_words)
+                encoded_images.append(image_to_base64(yolo_img))
 
             # if element_counts[element] > 1: #某隻角色的數量>1
-            #     output_path = f"{YOLO_DIR}/{message_id}/{element}"
-            #     output_file = get_jpg_files(output_path)
+            #     yolo_path = f"{YOLO_DIR}/{message_id}/{element}"
+            #     yolo_file = get_jpg_files(yolo_path)
 
-            #     for output_img in output_file: # 取得每張圖的路徑
-            #         encoded_images.append(image_to_base64(output_img))
+            #     for yolo_img in yolo_file: # 取得每張圖的路徑
+            #         encoded_images.append(image_to_base64(yolo_img))
 
             # else  : #某隻角色的數量=1
-            #     output_path = f"{YOLO_DIR}/{message_id}/{element}/im.jpg.jpg"
-            #     encoded_images.append(image_to_base64(output_path))
+            #     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,
-            'images': encoded_images,
             'description': element_list
+            'images': [
+                {
+                    'encoded_image': encoded_image,
+                    'description_list': top_k_words
+                }
+                for encoded_image, elements in zip(encoded_images, element_list)
+            ]
         }
+        # response_data = {
+        #     'message_id': message_id,
+        #     'images': encoded_images,
+        #     'description': element_list
+        # }
         
         return jsonify(response_data)
         
         # for label_name in label_names:
-        #     output_file=f"{YOLO_DIR}/{message_id}/{label_name}/im.jpg.jpg"
+        #     yolo_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')
+        #     encoded_images.append(image_to_base64(yolo_file))
 
     
     # # dictionary is not a JSON: https://www.quora.com/What-is-the-difference-between-JSON-and-a-dictionary