Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from deepface import DeepFace | |
| import cv2 | |
| import numpy as np | |
| import os | |
| import shutil | |
| # --- 設定 --- | |
| DB_PATH = "my_db" | |
| # 起動時にDBフォルダをリセット（Hugging Faceの環境をクリーンに保つため） | |
| if os.path.exists(DB_PATH): | |
| shutil.rmtree(DB_PATH) | |
| os.makedirs(DB_PATH, exist_ok=True) | |
| # --- 関数定義 --- | |
| def register_face(image, name): | |
| """画像をDBフォルダに保存する""" | |
| if image is None or name.strip() == "": | |
| return "画像と名前を入力してください。" | |
| try: | |
| # OpenCV形式に変換して保存 | |
| image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) | |
| # ここを修正：os.path.join です（スラッシュは不要） | |
| file_path = os.path.join(DB_PATH, f"{name}.jpg") | |
| cv2.imwrite(file_path, image_bgr) | |
| return f"「{name}」さんを登録しました！" | |
| except Exception as e: | |
| return f"エラーが発生しました: {str(e)}" | |
| def recognize_face(target_image): | |
| """画像から登録済みの顔を探す""" | |
| if target_image is None: | |
| return None, "画像をアップロードしてください。" | |
| if not os.listdir(DB_PATH): | |
| return target_image, "まだ誰も登録されていません。「推し登録」タブで登録してください。" | |
| try: | |
| # 認識の実行 | |
| results = DeepFace.find(img_path=target_image, db_path=DB_PATH, enforce_detection=False, silent=True) | |
| output_image = target_image.copy() | |
| found_names = [] | |
| for df in results: | |
| if df.empty: continue | |
| for index, row in df.iterrows(): | |
| # 登録名を取得 | |
| full_path = row['identity'] | |
| name = os.path.basename(full_path).split('.')[0] | |
| found_names.append(name) | |
| # 座標の取得と描画 | |
| x, y, w, h = int(row['source_x']), int(row['source_y']), int(row['source_w']), int(row['source_h']) | |
| cv2.rectangle(output_image, (x, y), (x+w, y+h), (0, 255, 0), 3) | |
| cv2.putText(output_image, name, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2) | |
| if not found_names: | |
| message = "登録された推しは見つかりませんでした。" | |
| else: | |
| message = f"発見しました！: {', '.join(set(found_names))}" | |
| return output_image, message | |
| except Exception as e: | |
| return target_image, f"認識エラー: {str(e)}" | |
| # --- UI (Gradio) の構築 --- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🌟 推し認識AIプロトタイプ") | |
| gr.Markdown("まずは『推し登録』で写真をアップし、その後『推しを探す』で判定してください。") | |
| with gr.Tabs(): | |
| with gr.TabItem("① 推し登録"): | |
| with gr.Row(): | |
| reg_input_img = gr.Image(label="推しの写真") | |
| reg_input_name = gr.Textbox(label="推しの名前") | |
| reg_btn = gr.Button("データベースに登録") | |
| reg_status = gr.Textbox(label="登録状況") | |
| reg_btn.click(register_face, inputs=[reg_input_img, reg_input_name], outputs=reg_status) | |
| with gr.TabItem("② 推しを探す"): | |
| with gr.Row(): | |
| rec_input_img = gr.Image(label="判定したい画像") | |
| rec_output_img = gr.Image(label="判定結果") | |
| rec_btn = gr.Button("推しを検索！") | |
| rec_status = gr.Textbox(label="結果メッセージ") | |
| rec_btn.click(recognize_face, inputs=[rec_input_img], outputs=[rec_output_img, rec_status]) | |
| # アプリ起動 | |
| if __name__ == "__main__": | |
| demo.launch() |