app.py

import gradio as gr
import cv2
import numpy as np
from PIL import Image
import qrcode
from googleapiclient.discovery import build
import os
import yaml
from ultralytics import YOLO
import logging
import glob
import re
import time  # 確保有導入 time 模組
import subprocess  # 用於執行外部命令
import threading  # 用於處理多線程
import queue  # 用於線程間通信
import urllib.request  # 用於檢查 URL 是否可訪問
import tempfile  # 用於創建臨時文件

markdown_content = """
<style>
.title-icon {
  display: inline-flex; /* 讓圖片和文字在同一行 */
  align-items: center; /* 垂直居中 */
  font-size: 2em; /* 調整整體大小 */
}
.title-icon img {
  height: 1em;  /* 圖片高度，與文字大小相關 */
  margin-right: 0.5em; /* 圖片與文字間距 */
}
</style>

<div class="title-icon">
  <h1>台灣常見鳥類偵測系統</h1>
</div>
"""

# 設定日誌
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)

# 載入類別名稱 (這部分保持不變)
try:
    logging.info("載入類別名稱...")
    yaml_path = "multibird_dataset.yaml"  # 改成你的路徑
    with open(yaml_path, "r", encoding="utf-8") as f:
        yaml_data = yaml.safe_load(f)
    class_names = yaml_data["names"]
    total_class_names = list(yaml_data["names"].values())
    logging.info(f"成功載入 {len(set(total_class_names))} 個類別")
except Exception as e:
    logging.error(f"載入類別名稱時發生錯誤: {str(e)}")
    class_names = {}

# 設定 YouTube API (這部分保持不變)
YOUTUBE_API_KEY = os.environ.get("YOUTUBE_API_KEY")
youtube = build("youtube", "v3", developerKey=YOUTUBE_API_KEY)

# 載入 YOLO 模型 (這部分保持不變)
try:
    logging.info("開始載入 YOLO 模型...")
    model_path = "best.pt"  # 0323 187種鳥的版本
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"找不到模型文件：{model_path}")
    model = YOLO(model_path)
    logging.info("YOLO 模型載入成功")
except Exception as e:
    logging.error(f"載入模型時發生錯誤：{str(e)}")
    raise

# 建立英文鳥名到中文鳥名的對照表
bird_chinese_names = {
    "domestic_goose": "家鵝",
    "black_swan": "黑天鵝",
    "spot_billed_duck": "花嘴鴨",
    "mallard": "綠頭鴨",
    "rock_dove": "野鴿",
    "spotted_dove": "珠頸斑鳩",
    "emerald_dove": "翠翼鳩",
    "common_moorhen": "紅冠水雞",
    "little_grebe": "小鸊鷉",
    "black_crowned_night_heron": "夜鷺",
    "malayan_night_heron": "黑冠麻鷺",
    "cattle_egret": "黃頭鷺",
    "mountain_hawk_eagle": "熊鷹",
    "common_kingfisher": "翠鳥",
    "muller_barbet": "五色鳥",
    "black_drongo": "大卷尾",
    "blue_flycatcher": "黑枕藍鶲",
    "formosan_blue_magpie": "台灣藍鵲",
    "chinese_bulbul": "白頭翁",
    "streak_breasted_scimitar_babbler": "小彎嘴",
    "black_necklaced_scimitar_babbler": "大彎嘴",
    "javan_myna": "家八哥",
    "white_vented_myna": "白腹八哥",
    "oriental_magpie_robin": "鵲鴝",
    "white_rumped_shama": "白腰鵲鴝",
    "daurian_redstart": "黃尾鴝",
    "eurasian_tree_sparrow": "麻雀",
    "white_wagtail": "白鶺鴒",
    "black_naped_oriole": "黃鸝",
    "maroon_oriole": "朱鸝",
    "black_faced_spoonbill": "黑面琵鷺",
    "great_egret": "大白鷺",
    "intermediate_egret": "中白鷺",
    "little_egret": "小白鷺",
    "chinese_egret": "唐白鷺",
    "grey_heron": "蒼鷺",
    "eurasian_spoonbill": "白琵鷺",
    "osprey": "魚鷹",
    "black_winged_kite": "黑翅鳶",
    "black_kite": "黑鳶",
    "mountain_eagle": "林鵰",
    "eastern_imperial_eagle": "花鵰",
    "oriental_honey_buzzard": "東方蜂鷹",
    "crested_serpent_eagle": "大冠鷲",
    "grey_faced_buzzard": "灰面鵟鷹",
    "crested_goshawk": "鳳頭蒼鷹",
    "chinese_sparrowhawk": "赤腹鷹",
    "besra": "松雀鷹",
    "white_tailed_eagle": "白尾海鵰",
    "peregrine_falcon": "遊隼",
    "ancient_murrelet": "扁嘴海雀",
    "tundra_swan": "小天鵝",
    "whooper_swan": "大天鵝",
    "great_crested_grebe": "鳳頭鸊鷉",
    "eared_grebe": "黑頸鸊鷉",
    "cotton_pygmy_goose": "棉鴨",
    "ferruginous_duck": "白眼潛鴨",
    "great_cormorant": "鸕鶿",
    "japanese_cormorant": "暗綠背鸕鶿",
    "brown_noddy": "褐翅燕鷗",
    "brown_headed_gull": "棕頭鷗",
    "black_headed_gull": "紅嘴鷗",
    "saunders_gull": "黑嘴鷗",
    "taiwan_hwamei": "台灣畫眉",
    "green_winged_teal": "小水鴨",
    "baikal_teal": "花臉鴨",
    "garganey": "白眉鴨",
    "common_shelduck": "翹鼻麻鴨",
    "common_pochard": "紅頭潛鴨",
    "anas_falcata": "羅紋鴨",
    "gadwall": "赤膀鴨",
    "mandarin_duck": "鴛鴦",
    "greater_scaup": "斑背潛鴨",
    "tufted_duck": "鳳頭潛鴨",
    "northern_pintail": "尖尾鴨",
    "eurasian_wigeon": "赤頸鴨",
    "northern_shoveller": "琵嘴鴨",
    "eastern_spot_billed_duck": "東方斑嘴鴨",
    "american_wigeon": "美洲赤頸鴨",
    "eurasian_coot": "白骨頂",
    "black_tailed_gull": "黑尾鷗",
    "caspian_gull": "裡海鷗",
    "caspian_tern": "紅嘴巨鷗",
    "relict_gull": "遺鷗",
    "graylag_goose": "灰雁",
    "swan_goose": "鴻雁",
    "tundra_bean_goose": "短嘴豆雁",
    "lesser_white_fronted_goose": "小白額雁",
    "taiga_bean_goose": "豆雁",
    "greater_white_fronted_goose": "大白額雁",
    "red_footed_booby": "紅腳鰹鳥",
    "franklin_gull": "富蘭克林鷗",
    "smew": "白秋沙",
    "brown_headed_thrush": "赤腹鶇",
    "grey_backed_thrush": "灰背鶇",
    "island_thrush": "島鶇",
    "plumbeous_redstart": "鉛色水鶇",
    "red_breasted_flycatcher": "紅胸姬鶲",
    "fairy_pitta": "八色鳥",
    "blue_winged_pitta": "藍翅八色鳥",
    "collared_kingfisher": "蒼翡翠",
    "ruddy_kingfisher": "赤翡翠",
    "dollar_bird": "三寶鳥",
    "eurasian_wryneck": "蟻鴷",
    "grey_capped_woodpecker": "小啄木",
    "japanese_waxwing": "日本太平鳥",
    "mikado_pheasant": "帝雉",
    "swinhoe_pheasant": "藍腹鷴",
    "ring_necked_pheasant": "環頸雉",
    "taiwan_bamboo_partridge": "台灣竹雞",
    "chestnut_backed_shrike": "栗背伯勞",
    "steppe_grey_strike": "草原灰伯勞",
    "brown_shrike": "紅尾伯勞",
    "long_tailed_shrike": "棕背伯勞",
    "taiwan_scimitar_babbler": "台灣彎嘴畫眉",
    "rufous_faced_warbler": "棕面鶯",
    "taiwan_yuhina": "冠羽畫眉",
    "green_backed_tit": "青背山雀",
    "yellow_tit": "黃山雀",
    "coal_tit": "煤山雀",
    "grey_chinned_minivet": "灰喉山椒鳥",
    "black_winged_cuckooshrike": "黑翅鵑鵙",
    "large_cuckoo_shrike": "大鵑鵙",
    "large_billed_crow": "巨嘴鴉",
    "olive_backed_pipit": "樹鷚",
    "alpine_accentor": "岩鷚",
    "eurasian_nuthatch": "茶腹鳾",
    "asian_koel": "噪鵑",
    "chestnut_bunting": "栗鵐",
    "little_bunting": "小鵐",
    "black_faced_bunting": "黑臉鵐",
    "yellow_throated_bunting": "黃喉鵐",
    "oriental_reed_warbler": "東方大葦鶯",
    "narcissus_flycatcher": "黃眉姬鶲",
    "mugimaki_flycatcher": "鴝姬鶲",
    "verditer_flycatcher": "銅藍鶲",
    "siberian_rubythroat": "紅喉歌鴝",
    "forest_wagtail": "山鶺鴒",
    "vinous_breasted_starling": "亞洲輝椋鳥",
    "plain_flowerpecker": "純色啄花鳥",
    "fire_breasted_flowerpecker": "紅胸啄花鳥",
    "taiwan_barbet": "台灣擬啄木",
    "collared_scops_owl": "領角鴞",
    "brown_hawk_owl": "褐鷹鴞",
    "brown_wood_owl": "褐林鴞",
    "collared_owlet": "鵂鶹",
    "austrialasian_grass_owl": "草鴞",
    "tawny_fish_owl": "黃魚鴞",
    "short_eared_owl": "短耳鴞",
    "hooded_crane": "白頭鶴",
    "siberian_crane": "白鶴",
    "sandhill_crane": "沙丘鶴",
    "common_crane": "灰鶴",
    "chinese_pond_heron": "池鷺",
    "pacific_reef_heron": "岩鷺",
    "purple_heron": "紫鷺",
    "yellow_bittern": "黃葦鳽",
    "great_bittern": "大麻鳽",
    "glossy_ibis": "彩鷺",
    "common_kestrel": "紅隼",
    "chinese_crested_tern": "中華鳳頭燕鷗",
    "white_eared_sibia": "白耳畫眉",
    "asian_dowitcher": "半蹼鷸",
    "bar_tailed_godwit": "斑尾塍鷸",
    "common": "田鷸",
    "far_eastern_curlew": "大杓鷸",
    "ruddy_turnstone": "翻石鷸",
    "little_stint": "小濱鷸",
    "curlew_sandpiper": "彎嘴濱鷸",
    "dunlin": "黑腹濱鷸",
    "spotted_redshank": "鶴鷸",
    "greater_painted_snipe": "彩鷸",
    "nordmann_greenshank": "諾氏鷸",
    "wood_sandpiper": "林鷸",
    "terek_sandpiper": "翹嘴鷸",
    "sanderling": "三趾濱鷸",
    "eurasian_curlew": "白腰杓鷸",
    "long_toed_stint": "長趾濱鷸",
    "common_greenshank": "青足鷸",
    "whimbrel": "中杓鷸",
    "marsh_sandpiper": "澤鷸",
    "red_necked_phalarope": "紅頸瓣足鷸",
    "common_redshank": "紅腳鷸",
    "little_ringed_plover": "小環頸鴴",
    "kentish_plover": "環頸鴴",
    "lesser_sand_plover": "蒙古鴴",
    "black_winged_stilt": "高蹺鴴",
    "oriental_pratincole": "燕鴴",
    "pacific_golden_plover": "太平洋金斑鴴",
    "pied_avocet": "反嘴鷸",
    "little_curlew": "小杓鷸",
    "common_ringed_plover": "劍鴴",
    "northern_lapwing": "田鳧",
    "grey_headed_lapwing": "灰頭麥雞",
    "savanna_nightjar": "南亞夜鷹",
    "little_tern": "小燕鷗",
    "glareola_maldivarum": "燕鴴",
    "barn_swallow": "家燕",
    "greater_crested_tern": "鳳頭燕鷗",
    "common_tern": "普通燕鷗",
    "whiskered_tern": "黑腹燕鷗",
    "gull_billed_tern": "鷗嘴燕鷗",
    "black_naped_tern": "黑枕燕鷗",
    "white_backed_woodpecker": "大赤啄木",
    "oriental_stork": "東方白鸛",
    "red_collared_dove": "紅鳩",
    "black_chinned_fruit_dove": "小綠鳩",
    "zebra_dove": "斑馬鳩",
    "common_goldeneye": "鵲鴨",
    "red_breasted_merganser": "紅胸秋沙",
    "scaly_sided_merganser": "中華秋沙",
    "great_white_pelican": "白鵜鶘",
    "bar_headed_goose": "斑頭雁",
    "adelie_penguin": "阿德利企鵝",
    "gentoo_penguin": "巴布亞企鵝",
    "chinstrap_penguin": "頰帶企鵝",
    "macaroni_penguin": "長冠企鵝",
    "king_penguin": "國王企鵝",
    "tufted_puffin": "簇絨海鸚",
    "hoopoe": "戴勝",
    "anna_hummingbird": "安氏蜂鳥",
}

# 將yaml中的類別名稱加入bird_chinese_names
for key, value in class_names.items():
    if key not in bird_chinese_names:
        bird_chinese_names[key] = value


def search_youtube_video(query):
    """搜尋 YouTube 影片"""
    try:
        # 將英文鳥名轉換為中文鳥名
        if query in bird_chinese_names:
            chinese_query = bird_chinese_names[query]
            logging.info(f"將英文鳥名 '{query}' 轉換為中文鳥名 '{chinese_query}'")
            query = chinese_query
        else:
            logging.warning(f"未找到 '{query}' 的中文對照名稱，使用原始名稱進行搜尋")

        logging.info(f"搜尋 YouTube 影片，關鍵字：{query}")
        search_response = (
            youtube.search()
            .list(
                q=query,
                part="id,snippet",
                maxResults=1,
                type="video",
                channelId="UClCaVtB-pYpNOJGiGzQAhNA",  # 你的頻道 ID
            )
            .execute()
        )

        if not search_response["items"]:
            logging.warning("未找到相關影片")
            return None, None, None

        video_id = search_response["items"][0]["id"]["videoId"]
        video_url = f"https://www.youtube.com/watch?v={video_id}"
        logging.info(f"找到影片：{video_url}")

        embed_html = f'<iframe width="100%" height="315" src="https://www.youtube.com/embed/{video_id}" frameborder="0" allowfullscreen></iframe>'

        # 生成 QR Code
        qr = qrcode.QRCode(version=1, box_size=10, border=5)
        qr.add_data(video_url)
        qr.make(fit=True)
        qr_img = qr.make_image(fill_color="black", back_color="white")
        qr_array = np.array(qr_img.convert("RGB"))  # 轉換為RGB模式以確保兼容性

        return video_url, embed_html, qr_array
    except Exception as e:
        logging.error(f"YouTube 搜尋錯誤: {str(e)}")
        return None, None, None


def is_valid_youtube_url(url):
    """檢查是否為有效的 YouTube URL"""
    youtube_regex = (
        r"(https?://)?(www\.)?"
        r"(youtube|youtu|youtube-nocookie)\.(com|be)/"
        r"(watch\?v=|embed/|v/|.+\?v=)?([^&=%\?]{11})"
    )
    match = re.match(youtube_regex, url)
    return match is not None


def get_youtube_video_id(url):
    """從 YouTube URL 中提取視頻 ID"""
    youtube_regex = (
        r"(https?://)?(www\.)?"
        r"(youtube|youtu|youtube-nocookie)\.(com|be)/"
        r"(watch\?v=|embed/|v/|.+\?v=)?([^&=%\?]{11})"
    )
    match = re.match(youtube_regex, url)
    if match:
        return match.group(6)
    return None


def process_image(input_image):
    """處理圖片並進行物件偵測 (這部分保持不變)"""
    try:
        if input_image is None:
            logging.warning("未提供輸入圖片")
            return None, "請先上傳圖片", "", None

        logging.info("開始處理圖片...")
        # 建立輸出目錄
        output_dir = "results"
        os.makedirs(output_dir, exist_ok=True)

        # 執行物件偵測
        results = model.predict(input_image, save=False)
        result = results[0]

        # 取得偵測結果
        detected_objects = []
        for box in result.boxes:
            cls = int(box.cls[0])
            conf = float(box.conf[0])
            if conf > 0.8:  # 信心度閾值
                label = class_names.get(cls, f"類別 {cls}")
                detected_objects.append(label)
                logging.info(f"檢測到物件：{label}，信心度：{conf:.2f}")

        # 取得處理後的圖片
        result_image = result.plot()
        result_pil = Image.fromarray(result_image)

        # 搜尋布萊克獨木舟生態休閒中心 YouTube 影片
        if detected_objects:
            video_url, embed_html, qr_image = search_youtube_video(detected_objects[0])
            result_text = f"檢測到: {', '.join(detected_objects)}"
        else:
            video_url = "偵測該物件的信心度不足，未查詢到相關影片，偵測結果僅供參考"
            embed_html = ""
            qr_image = None
            result_text = "偵測該物件的信心度不足，未查詢到相關影片，偵測結果僅供參考"

        return result_pil, video_url, embed_html, qr_image
    except Exception as e:
        logging.error(f"處理錯誤: {str(e)}")
        return None, f"處理錯誤: {str(e)}", "", None


def process_video(input_video):
    """處理影片並進行物件偵測"""
    try:
        if input_video is None:
            logging.warning("未提供輸入影片")
            return None, "未提供輸入影片", "", None  # 修正返回的訊息數量

        logging.info("開始處理影片...")
        output_dir = "results"
        os.makedirs(output_dir, exist_ok=True)

        timestamp = int(time.time())
        run_name = f"predict_{timestamp}"
        logging.info(f"使用唯一識別名稱進行偵測: {run_name}")

        # 執行物件偵測，並直接處理結果
        results = model(
            input_video, save=True, project=output_dir, name=run_name, stream=True
        )

        # 儲存處理過的frame
        processed_frames = []

        # 取得偵測結果和處理frame
        detected_objects = set()
        for result in results:
            for box in result.boxes:
                cls = int(box.cls[0])
                conf = float(box.conf[0])
                if conf > 0.8:
                    label = class_names.get(cls, f"類別 {cls}")
                    detected_objects.add(label)
                    logging.info(f"檢測到物件：{label}，信心度：{conf:.2f}")
            # 將result.plot()存入processed_frames
            processed_frames.append(result.plot())

        # 影片格式轉換和輸出
        predict_folder = os.path.join(output_dir, run_name)
        output_video_name = os.path.join(predict_folder, f"output_{timestamp}.mp4")

        if not os.path.exists(predict_folder):
            os.makedirs(predict_folder)

        # 取得影片的基本資訊 (從原始影片取得)
        cap = cv2.VideoCapture(input_video)  # 改為讀取原始影片
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        cap.release()  # 讀取完後就釋放

        # 建立 VideoWriter 物件 (使用 mp4v codec)
        fourcc = cv2.VideoWriter_fourcc(*"avc1")
        out = cv2.VideoWriter(output_video_name, fourcc, fps, (width, height))

        # 將處理過的 frame 寫入影片
        for frame in processed_frames:
            out.write(frame)
        out.release()

        # 搜尋 YouTube 影片 (其餘部分保持不變)
        if detected_objects:
            video_url, embed_html, qr_image = search_youtube_video(
                list(detected_objects)[0]
            )
            result_text = f"檢測到: {', '.join(detected_objects)}"
        else:
            video_url = "偵測該物件的信心度不足，未查詢到相關影片，偵測結果僅供參考"
            embed_html = ""
            qr_image = None
            result_text = "偵測該物件的信心度不足，未查詢到相關影片，偵測結果僅供參考"

        return output_video_name, video_url, embed_html, qr_image  # 直接回傳名稱

    except Exception as e:
        logging.error(f"影片處理錯誤: {str(e)}")
        return None, f"影片處理錯誤: {str(e)}", "", None


def search_channel_videos(query):
    """搜尋布萊克獨木舟生態休閒中心頻道的影片"""
    try:
        logging.info(f"搜尋頻道影片，關鍵字：{query}")
        search_response = (
            youtube.search()
            .list(
                q=query,
                part="id,snippet",
                maxResults=6,
                type="video",
                channelId="UClCaVtB-pYpNOJGiGzQAhNA",  # 布萊克獨木舟生態休閒中心頻道 ID
                order="relevance"
            )
            .execute()
        )

        results = []
        for item in search_response.get("items", []):
            video_id = item["id"]["videoId"]
            title = item["snippet"]["title"]
            description = item["snippet"]["description"]
            published_at = item["snippet"]["publishedAt"]
            embed_html = f'<div style="margin-bottom: 20px;"><h3>{title}</h3><p>發布時間：{published_at}</p><iframe width="100%" height="315" src="https://www.youtube.com/embed/{video_id}" frameborder="0" allowfullscreen></iframe><p>{description}</p></div>'
            results.append(embed_html)

        if not results:
            return "<p>查無相關影片</p>"
        return "<div style='max-width: 800px; margin: 0 auto;'>" + "".join(results) + "</div>"
    except Exception as e:
        logging.error(f"頻道影片搜尋錯誤: {str(e)}")
        return f"<p>搜尋錯誤: {str(e)}</p>"


# 創建 Gradio 介面
with gr.Blocks(title="鳥類偵測系統") as demo:
    gr.Markdown(markdown_content)
    gr.Markdown("### 支援圖片、影片、即時攝影機")

    with gr.Tab("圖片偵測"):
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(label="上傳圖片", type="numpy")
                image_button = gr.Button("開始偵測")
            with gr.Column(scale=1):
                image_output = gr.Image(label="偵測結果")
        with gr.Row():
            with gr.Column(scale=1):
                image_video_url = gr.Textbox(
                    label="YouTube 影片連結", visible=True
                )  # 保留但隱藏
                image_video_embed = gr.HTML(label="YouTube 影片")
            with gr.Column(scale=1):
                image_qr_output = gr.Image(label="QR Code")

        image_button.click(
            process_image,
            inputs=[image_input],
            outputs=[image_output, image_video_url, image_video_embed, image_qr_output],
        )

        # 新增頻道影片搜尋頁籤
    with gr.Tab("頻道影片搜尋"):
        gr.Markdown("### 布萊克獨木舟生態休閒中心頻道影片搜尋")
        with gr.Row():
            search_input = gr.Textbox(
                label="請輸入關鍵字",
                placeholder="輸入要搜尋的關鍵字...",
                scale=4
            )
            search_button = gr.Button("搜尋", scale=1)

        with gr.Row():
            results_html = gr.HTML(
                label="搜尋結果",
                value="<p>請輸入關鍵字進行搜尋</p>"
            )

        search_button.click(
            fn=search_channel_videos,
            inputs=[search_input],
            outputs=[results_html]
        )

    # 影片偵測頁籤
    # with gr.Tab("影片偵測"):
    #     with gr.Row():
    #         with gr.Column(scale=1):
    #             video_input = gr.Video(label="上傳影片")
    #             video_button = gr.Button("開始偵測")
    #         with gr.Column(scale=1):
    #             video_output = gr.Video(
    #                 label="偵測結果",
    #                 # streaming=True,  # 不再需要 streaming
    #                 # autoplay=True   # autoplay 根據你的需求決定是否保留
    #             )
    #     with gr.Row():
    #         with gr.Column(scale=1):
    #             video_video_url = gr.Textbox(
    #                 label="YouTube 影片連結", visible=True
    #             )
    #             video_video_embed = gr.HTML(label="YouTube 影片")
    #         with gr.Column(scale=1):
    #             video_qr_output = gr.Image(label="QR Code")

    #     video_button.click(
    #         process_video,
    #         inputs=[video_input],
    #         outputs=[video_output, video_video_url, video_video_embed, video_qr_output],
    #     )

# 啟動應用程式
if __name__ == "__main__":
    logging.info("啟動網頁應用程式...")
    demo.launch(share=True, server_name="0.0.0.0", pwa=True)