Update main.py

main.py

@@ -12,14 +12,15 @@ import uvicorn
 import logging
 from PIL import Image
 import io
+import requests.exceptions
 
 # 設置日誌記錄
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # 檢查環境變數
-if not all(k in os.environ for k in ["CHANNEL_ACCESS_TOKEN", "CHANNEL_SECRET", "GOOGLE_API_KEY", "IMGBB_API_KEY"]):
-    logging.error("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, and IMGBB_API_KEY.")
-    raise ValueError("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, and IMGBB_API_KEY.")
+if not all(k in os.environ for k in ["CHANNEL_ACCESS_TOKEN", "CHANNEL_SECRET", "GOOGLE_API_KEY", "IMGBB_API_KEY", "HF_SPACE"]):
+    logging.error("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, IMGBB_API_KEY, and HF_SPACE.")
+    raise ValueError("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, IMGBB_API_KEY, and HF_SPACE.")
 
 # Line Bot API 設定
 line_bot_api = LineBotApi(os.environ["CHANNEL_ACCESS_TOKEN"])
@@ -42,6 +43,7 @@ GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini
 IMAGIN_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent"
 GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]
 IMGBB_API_KEY = os.environ["IMGBB_API_KEY"]
+HF_SPACE_URL = os.environ["HF_SPACE"]
 
 app = FastAPI()
 
@@ -139,7 +141,6 @@ def get_gemini_response(prompt: str, images: list, api_url: str):
         logging.error(f"Gemini API request failed: {e}")
         return f"Gemini API 請求失敗：{e}"
 
-
 def get_virtual_tryon_image(user_photo_base64: str, upper_body_image_base64: str, lower_body_image_base64: str):
     """
     使用 gemini-2.5-flash-image-preview (nanobanana) 模擬虛擬試穿。
@@ -192,6 +193,81 @@ def get_virtual_tryon_image(user_photo_base64: str, upper_body_image_base64: str
         logging.error(f"Gemini image generation API request failed: {e}")
         return None
 
+def generate_outfit_image(upper_body_images: list, lower_body_images: list):
+    """
+    使用 gemini-2.5-flash-image-preview 生成一套全新的穿搭圖片。
+    """
+    logging.info("Attempting to generate a new outfit image.")
+
+    prompt = "使用提供的上衣和褲子/裙子圖片，生成一套完整且時尚的穿搭圖片。請將衣服呈現在一個有模特兒穿著或是在平面上呈現的完整畫面中。風格應與提供的衣物相符。"
+    
+    all_images_base64 = (
+        [get_base64_from_url(url) for url in upper_body_images] +
+        [get_base64_from_url(url) for url in lower_body_images]
+    )
+
+    payload = {
+        "contents": [
+            {
+                "parts": [
+                    {"text": prompt}
+                ] + [
+                    {"inlineData": {"mimeType": "image/jpeg", "data": img_base64}} for img_base64 in all_images_base64
+                ]
+            }
+        ],
+        "generationConfig": {
+            "responseModalities": ['IMAGE']
+        }
+    }
+
+    try:
+        response = requests.post(f"{IMAGIN_API_URL}?key={GOOGLE_API_KEY}", json=payload, timeout=120)
+        response.raise_for_status()
+        response_data = response.json()
+        if 'candidates' in response_data and response_data['candidates']:
+            image_part = response_data['candidates'][0]['content']['parts'][0]
+            if image_part and 'inlineData' in image_part:
+                generated_image_base64 = image_part['inlineData']['data']
+                generated_image_url = upload_to_imgbb(generated_image_base64)
+                if generated_image_url:
+                    logging.info(f"Successfully generated and uploaded new outfit image: {generated_image_url}")
+                    return generated_image_url
+                else:
+                    logging.error("Failed to upload generated outfit image to Imgbb.")
+                    return None
+            else:
+                logging.error("Gemini image response format is invalid for outfit generation.")
+                return None
+        else:
+            logging.error(f"Gemini outfit generation response has no candidates: {response.text}")
+            return None
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Gemini outfit generation API request failed: {e}")
+        return None
+
+def save_image_locally(image_binary: bytes):
+    """
+    將二進位圖片資料儲存到本地，並返回一個可供外部存取的 URL。
+    """
+    try:
+        # 確保 'static' 資料夾存在
+        if not os.path.exists("static"):
+            os.makedirs("static")
+        
+        image = Image.open(io.BytesIO(image_binary))
+        # 隨機生成一個檔案名以避免衝突
+        file_name = f"static/{os.urandom(16).hex()}.png"
+        image.save(file_name, format="PNG")
+        
+        image_url = os.path.join(HF_SPACE_URL, file_name)
+        logging.info(f"Image successfully saved locally: {image_url}")
+        return image_url
+    except Exception as e:
+        logging.error(f"Error saving image locally: {e}")
+        return None
+
+
 @line_handler.add(MessageEvent, message=TextMessage)
 def handle_text_message(event):
     user_id = event.source.user_id
@@ -209,8 +285,46 @@ def handle_text_message(event):
         )
         return
     
+    # 新增 圖片生成穿搭功能
+    if text in ["生成穿搭", "圖片推薦"] and user_states[user_id]["is_ready_for_outfit"]:
+        if len(user_states[user_id]["upper_body_images"]) == MAX_IMAGES_PER_TYPE and len(user_states[user_id]["lower_body_images"]) == MAX_IMAGES_PER_TYPE:
+            line_bot_api.reply_message(
+                reply_token,
+                TextSendMessage(text="好的，正在為您生成一套新的穿搭圖片，請稍候...")
+            )
+            
+            generated_outfit_url = generate_outfit_image(
+                user_states[user_id]["upper_body_images"],
+                user_states[user_id]["lower_body_images"]
+            )
+
+            if generated_outfit_url:
+                line_bot_api.push_message(
+                    user_id,
+                    ImageSendMessage(original_content_url=generated_outfit_url, preview_image_url=generated_outfit_url)
+                )
+                line_bot_api.push_message(
+                    user_id,
+                    TextSendMessage(text="這是根據您的衣物生成的圖片推薦。如果想再次使用，請輸入「重置」。")
+                )
+            else:
+                line_bot_api.push_message(
+                    user_id,
+                    TextSendMessage(text="圖片生成失敗，請稍後再試。")
+                )
+            
+            # 生成後重置狀態
+            user_states[user_id] = defaultdict(lambda: {"upper_body_images": [], "lower_body_images": [], "current_mode": None, "is_ready_for_outfit": False, "is_ready_for_photo": False, "user_info": {}, "personal_photo": None, "personal_photo_base64": None})[user_id]
+            return
+        else:
+            line_bot_api.reply_message(
+                reply_token,
+                TextSendMessage(text="請先上傳三件上衣和三件褲子/裙子圖片。")
+            )
+            return
+
     # 處理個人資訊輸入
-    if not user_states[user_id]["is_ready_for_photo"] and not user_states[user_id]["is_ready_for_outfit"]:
+    if not user_states[user_id]["is_ready_for_outfit"] and not user_states[user_id]["is_ready_for_photo"]:
         try:
             parts = text.split(',')
             if len(parts) == 5:
@@ -238,7 +352,7 @@ def handle_text_message(event):
             logging.error(f"Invalid user info format: '{text}'")
             line_bot_api.reply_message(
                 reply_token,
-                TextSendMessage(text="請依照格式輸入：身高,三圍胸圍,腰圍,臀圍,場合。例如：165,85,65,90,約會")
+                TextSendMessage(text="請依照格式輸入：身高,胸圍,腰圍,臀圍,場合。例如：165,85,65,90,約會")
             )
             return
 
@@ -265,6 +379,7 @@ def handle_text_message(event):
         TextSendMessage(text="請先依照格式輸入個人資訊，格式為：身高,胸圍,腰圍,臀圍,場合。例如：165,85,65,90,約會")
     )
 
+
 @line_handler.add(MessageEvent, message=ImageMessage)
 def handle_image_message(event):
     user_id = event.source.user_id
@@ -281,17 +396,24 @@ def handle_image_message(event):
         )
         return
 
+    # 檢查是否所有衣物都已上傳
+    is_outfit_complete = (
+        len(user_states[user_id]["upper_body_images"]) == MAX_IMAGES_PER_TYPE and
+        len(user_states[user_id]["lower_body_images"]) == MAX_IMAGES_PER_TYPE
+    )
+
     # 處理個人照片上傳
-    if len(user_states[user_id]["upper_body_images"]) == MAX_IMAGES_PER_TYPE and len(user_states[user_id]["lower_body_images"]) == MAX_IMAGES_PER_TYPE and not user_states[user_id]["personal_photo"]:
+    if is_outfit_complete and not user_states[user_id]["personal_photo"]:
         try:
             image_id = event.message.id
-            base64_img = get_base64_image(image_id)
-            
-            # 先將個人照片的 Base64 儲存起來，用於 Gemini
-            user_states[user_id]["personal_photo_base64"] = base64_img
+            # 取得圖片的二進位資料
+            image_content = line_bot_api.get_message_content(image_id).content
             
-            # 再將個人照片上傳到 Imgbb，並儲存 URL
-            photo_url = upload_to_imgbb(base64_img)
+            # 使用新函式來儲存圖片並取得 URL
+            photo_url = save_image_locally(image_content)
+            if not photo_url:
+                raise Exception("Failed to save personal photo locally.")
+
             user_states[user_id]["personal_photo"] = photo_url
             user_states[user_id]["is_ready_for_photo"] = True
             
@@ -304,17 +426,17 @@ def handle_image_message(event):
             user_info = user_states[user_id]["user_info"]
             occasion = user_info["occasion"]
             
-            # Gemini 提示詞現在使用圖片 URL，而不是 Base64 編碼
             prompt = (
                 f"我提供了三件上衣圖片和三件下半身圖片，以及使用者的一張個人照片。使用者資訊：身高 {user_info['height']}cm，三圍 {user_info['bust']}-{user_info['waist']}-{user_info['hip']}，場合是「{occasion}」。"
                 "請根據這些衣物，為我推薦一套最適合的穿搭，並詳細說明為何這套搭配適合這個場合。請以繁體中文回答。"
                 "請在描述中以圖片連結的形式（例如：'https://example.com/outfit1.jpg'）顯示你搭配好的衣物組合。"
             )
             
+            # 取得圖片的 Base64 編碼，用於 Gemini API 呼叫
             all_images_base64 = (
                 [get_base64_from_url(url) for url in user_states[user_id]["upper_body_images"]] +
                 [get_base64_from_url(url) for url in user_states[user_id]["lower_body_images"]] +
-                [user_states[user_id]["personal_photo_base64"]]
+                [get_base64_from_url(photo_url)]
             )
             
             response_text = get_gemini_response(prompt, all_images_base64, GEMINI_API_URL)
@@ -328,7 +450,7 @@ def handle_image_message(event):
             best_lower_body_image = user_states[user_id]["lower_body_images"][best_lower_index]
 
             # 呼叫虛擬試穿 API，並取得結果圖片 URL
-            virtual_tryon_url = get_virtual_tryon_image(user_states[user_id]["personal_photo_base64"], get_base64_from_url(best_upper_body_image), get_base64_from_url(best_lower_body_image))
+            virtual_tryon_url = get_virtual_tryon_image(get_base64_from_url(photo_url), get_base64_from_url(best_upper_body_image), get_base64_from_url(best_lower_body_image))
             
             # 發送 Gemini 的文字建議
             line_bot_api.push_message(
@@ -367,11 +489,12 @@ def handle_image_message(event):
 
     try:
         image_id = event.message.id
-        base64_img = get_base64_image(image_id)
-        # 上傳圖片到 Imgbb，並儲存 URL
-        image_url = upload_to_imgbb(base64_img)
+        image_content = line_bot_api.get_message_content(image_id).content
+        
+        # 使用新函式來儲存圖片並取得 URL
+        image_url = save_image_locally(image_content)
         if not image_url:
-            logging.error("Failed to upload image to Imgbb.")
+            logging.error("Failed to save image locally.")
             line_bot_api.reply_message(
                 reply_token,
                 TextSendMessage(text="圖片上傳失敗，請稍後再試。")
@@ -381,9 +504,12 @@ def handle_image_message(event):
         if mode == "upper":
             if len(user_states[user_id]["upper_body_images"]) < MAX_IMAGES_PER_TYPE:
                 user_states[user_id]["upper_body_images"].append(image_url)
+                reply_text = f"已接收上衣。上衣: {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE}。"
+                if len(user_states[user_id]["upper_body_images"]) == MAX_IMAGES_PER_TYPE:
+                    reply_text += "上衣已收集完畢。"
                 line_bot_api.reply_message(
                     reply_token,
-                    TextSendMessage(text=f"已接收上衣。上衣: {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE}。")
+                    TextSendMessage(text=reply_text)
                 )
                 return
             else:
@@ -396,10 +522,20 @@ def handle_image_message(event):
         else: # mode == "lower"
             if len(user_states[user_id]["lower_body_images"]) < MAX_IMAGES_PER_TYPE:
                 user_states[user_id]["lower_body_images"].append(image_url)
+                reply_text = f"已接收褲子/裙子。上衣: {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE}。"
+                if len(user_states[user_id]["lower_body_images"]) == MAX_IMAGES_PER_TYPE:
+                    reply_text += "褲子/裙子已收集完畢。"
                 line_bot_api.reply_message(
                     reply_token,
-                    TextSendMessage(text=f"已接收褲子/裙子。上衣: {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE}。")
+                    TextSendMessage(text=reply_text)
                 )
+                # 檢查是否所有圖片都已上傳完畢
+                if (len(user_states[user_id]["upper_body_images"]) == MAX_IMAGES_PER_TYPE and
+                    len(user_states[user_id]["lower_body_images"]) == MAX_IMAGES_PER_TYPE):
+                    line_bot_api.push_message(
+                        user_id,
+                        TextSendMessage(text="所有衣物圖片已收集完畢！\n\n現在您可以上傳**個人照片**進行虛擬試穿，或輸入「**生成穿搭**」來獲得圖片推薦。")
+                    )
                 return
             else:
                 logging.info("Lower body image count is full.")
@@ -417,4 +553,4 @@ def handle_image_message(event):
         )
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8080)))
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8080)))