Update api_server.py

api_server.py

@@ -15,6 +15,7 @@ import torch
 from collections import Counter
 import psutil
 from gradio_client import Client, handle_file
+from io import BytesIO
 
 # Disable tensorflow warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
@@ -24,7 +25,8 @@ load_type = 'local'
 MODEL_YOLO = "yolo11_detect_best_241024_1.pt"
 MODEL_DIR = "./artifacts/models"
 YOLO_DIR = "./artifacts/yolo"
-GRADIO_URL = "https://fd39e54bcb191a37bf.gradio.live/"
+IMG2TEXT_URL = "https://fd39e54bcb191a37bf.gradio.live/"
+TEXT2IMG_URL = "https://91698ded8ba92d3bb0.gradio.live/"
 
 
 # Load the saved YOLO model into memory
@@ -36,7 +38,7 @@ if load_type == 'local':
         
     model = YOLO(model_path)
 
-    print("***** 1. LOAD YOLO MODEL DONE *****")
+    print("***** FLASK API---LOAD YOLO MODEL DONE *****")
     #model.eval()  # 設定模型為推理模式
 elif load_type == 'remote_hub_download':
     from huggingface_hub import hf_hub_download
@@ -62,6 +64,18 @@ def image_to_base64(image_path):
         encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
         return encoded_string
 
+def convert_webp_to_base64(webp_path):
+    # 開啟 .webp 圖片檔
+    with Image.open(webp_path) as img:
+        # 將圖片存到 BytesIO 物件中，以便轉換為 base64
+        buffered = BytesIO()
+        img.save(buffered, format="WEBP")
+        
+        # 取得 base64 編碼的字串
+        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        
+        return img_base64
+
 
 # 抓取指定路徑下的所有 JPG 檔案
 def get_jpg_files(path):
@@ -117,11 +131,11 @@ def predict():
     except Exception as e:
         return jsonify({'error': str(e)}), 400
         
-    print("***** 2. Start YOLO predict *****")
+    print("***** FLASK API---/predict Start YOLO predict *****")
     # Make a prediction using YOLO
     results = model(image_data)
-    print ("===== YOLO predict result:",results,"=====")
-    print("***** YOLO predict DONE *****")
+    print ("===== FLASK API---/predict YOLO predict result:",results,"=====")
+    print("***** FLASK API---/predict YOLO predict DONE *****")
     
     check_memory_usage()
     
@@ -145,7 +159,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"====== 3. YOLO label_names: {label_names}======")
+        print(f"====== FLASK API---/predict 3. YOLO label_names: {label_names}======")
 
         element_counts = Counter(label_names)
 
@@ -154,15 +168,15 @@ def predict():
             yolo_path = f"{YOLO_DIR}/{message_id}/{element}"
             yolo_file = get_jpg_files(yolo_path)
 
-            print(f"***** 處理:{yolo_path} *****")
+            print(f"***** FLASK API---/predict 處理:{yolo_path} *****")
 
             if len(yolo_file) == 0:
-                print(f"警告：{element} 沒有找到相關的 JPG 檔案")
+                print(f" FLASK API---/predict 警告：{element} 沒有找到相關的 JPG 檔案")
                 continue
 
             for yolo_img in yolo_file: # 每張切圖yolo_img
-                print("***** 4. START CLIP *****")
-                client = Client(GRADIO_URL)
+                print("***** FLASK API---/predict 4. START CLIP *****")
+                client = Client(IMG2TEXT_URL)
                 clip_result = client.predict(
                     image=handle_file(yolo_img),
                     top_k=3,
@@ -171,7 +185,7 @@ def predict():
                 top_k_words.append(clip_result) # CLIP預測3個結果(top_k_words)
                 encoded_images.append(image_to_base64(yolo_img))
                 element_list.append(element)
-                print(f"===== CLIP result:{top_k_words} =====\n")
+                print(f"===== FLASK API---/predict CLIP result:{top_k_words} =====\n")
                 
         # 建立回應資料
         response_data = {
@@ -193,14 +207,28 @@ def predict():
 
 
 # API route for health check
-@app.route('/health', methods=['GET'])
-def health():
-    """
-    Health check API to ensure the application is running.
-    Returns "OK" if the application is healthy.
-    Demo Usage: "curl http://localhost:5000/health" or using alias "curl http://127.0.0.1:5000/health"
-    """
-    return 'OK'
+@app.route('/text2img', methods=['POST'])
+def text2img():
+    text_message = request.form.get('text_message')
+    message_id = request.form.get('message_id')
+
+    client = Client(TEXT2IMG_URL)
+    result = client.predict(
+    		word= text_message,
+    		api_name="/predict"
+    )
+    print(f"===== FLASK API---/text2img 文字轉圖片result[0]:{result[0]} =====")
+    result_img = convert_webp_to_base64(result[0])
+    print(f"===== FLASK API---/text2img 文字轉圖片轉base64:{result_img} =====")
+
+    # 建立回應資料
+    response_data = {
+        'message_id': message_id,
+        'encoded_image': result_img,
+        'description': result[1]        
+    }
+
+    return jsonify(response_data), 200 
 
 
 # API route for version