Update api_server.py

api_server.py

@@ -25,8 +25,7 @@ load_type = 'local'
 MODEL_YOLO = "yolo11_detect_best_241024_1.pt"
 MODEL_DIR = "./artifacts/models"
 YOLO_DIR = "./artifacts/yolo"
-IMG2TEXT_URL = "https://f2d87a93eed6d95942.gradio.live/" # 圖生文
-TEXT2IMG_URL = "https://1fce6478b5a17e2fda.gradio.live/" # 文生圖
+GRADIO_URL = "https://f1a7e15aa545f4e172.gradio.live/" 
 
 
 # Load the saved YOLO model into memory
@@ -64,18 +63,6 @@ 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):
@@ -92,6 +79,19 @@ def get_jpg_files(path):
 # jpg_files = get_jpg_files(image_folder)
 
 
+def clip_model (choice="find_similar_words",image=None,word=None):
+    client = Client(GRADIO_URL)
+    clip_result = client.predict(
+		choice = choice,
+		image= handle_file(image),
+		word = word,
+		top_k = 3,
+		api_name = "/run_function"
+    )
+    return clip_result
+
+
+
 def check_memory_usage():
     # Get memory details
     memory_info = psutil.virtual_memory()
@@ -120,6 +120,8 @@ def predict():
     #user_id = request.args.get('user_id') 
     file = request.files['image']
     message_id = request.form.get('message_id') #str(uuid.uuid4())
+    choice = request.form.get('choice')
+    word = request.form.get('word')
     
     if 'image' not in request.files:
         # Handle if no file is selected
@@ -132,8 +134,10 @@ def predict():
         return jsonify({'error': str(e)}), 400
         
     print("***** FLASK API---/predict Start YOLO predict *****")
+    
     # Make a prediction using YOLO
     results = model(image_data)
+    
     print ("===== FLASK API---/predict YOLO predict result:",results,"=====")
     print("***** FLASK API---/predict YOLO predict DONE *****")
     
@@ -176,13 +180,10 @@ def predict():
 
             for yolo_img in yolo_file: # 每張切圖yolo_img
                 print("***** FLASK API---/predict 4. START CLIP *****")
-                client = Client(IMG2TEXT_URL)
-                clip_result = client.predict(
-                    image=handle_file(yolo_img),
-                    top_k=3,
-                    api_name="/predict"
-                )
-                top_k_words.append(clip_result) # CLIP預測3個結果(top_k_words)
+                
+                clip_model(choice,yolo_img,word)
+                
+                top_k_words.append(clip_result[0]) # CLIP預測3個結果(top_k_words)
                 encoded_images.append(image_to_base64(yolo_img))
                 element_list.append(element)
                 print(f"===== FLASK API---/predict CLIP RESULT:{top_k_words} =====\n")
@@ -213,22 +214,20 @@ def predict():
 # API route for health check
 @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:{result} =====")
-    result_img = result[2] # 已經是base64 coded
+    choice = request.form.get('choice')
+    word = request.form.get('word')
+    
+    clip_model(choice,None,word)
+    
+    print(f"===== FLASK API---/text2img 文字轉圖片result:{clip_result} =====")
+    result_img = clip_result[2] # 已經是base64 coded
 
     # 建立回應資料
     response_data = {
         'message_id': message_id,
         'encoded_image': result_img,
-        'description': result[1]        
+        'description': clip_result[0]        
     }
 
     return jsonify(response_data), 200