add

combined_data/blendshape_annotation_preprocess.py

@@ -0,0 +1,114 @@
+import os
+import json
+import torch
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from loguru import logger
+from typing import Dict, List, Tuple
+
+# ---- あなたの定義済み Dataset, transforms があるなら import ----
+# from your_dataset import BlendShapeDataset, image_transform
+
+def unify_index(idx: int, group_size: int) -> int:
+    """
+    BlendShape の選択 index が group_size と同じ場合、-1 (none) に変換して返す。
+    """
+    if idx == group_size:
+        return -1
+    return idx
+
+def build_combination_to_filename(
+    teacher_data_file: str,
+    meta_file: str,
+) -> Tuple[Dict[Tuple[int, ...], str], List[int]]:
+    """
+    BlendShapeData.json を読み込み、
+    例:
+       { (g0_idx, g1_idx, g2_idx): "000001.png", ... }
+    のような辞書を作成して返す。
+
+    さらに各グループのサイズ (blendShapeNames数 + 1) のリスト group_sizes も返す。
+    """
+    with open(meta_file, "r", encoding="utf-8") as f:
+        meta = json.load(f)
+    blend_shape_groups = meta["blendShapeGroupsMeta"]
+    # group_sizes[i] = len(そのグループの blendShapeNames) + 1(none枠)
+    group_sizes = [len(g["blendShapeNames"]) + 1 for g in blend_shape_groups]
+
+    # teacher_data_file 読み込み
+    with open(teacher_data_file, "r", encoding="utf-8") as f:
+        teacher_data = json.load(f)
+    data_list = teacher_data["dataList"]
+
+    combination_to_filename = {}
+    for data in data_list:
+        photo_filename = data["photoFileName"]
+        blendShapeSelections = data["blendShapeSelectionsPerGroup"]
+
+        # グループ順に selectedBlendShapeIndex を取得しつつ、-1の場合は-1のまま、
+        # group_sizeと一致していたら-1へ変換(理想的には -1 しか登場しない想定だが一応対応)
+        combo = []
+        for group_idx, selection in enumerate(blendShapeSelections):
+            sel_idx = selection["selectedBlendShapeIndex"]
+            # group_sizes[group_idx] と同じなら none として -1
+            sel_idx = unify_index(sel_idx, group_sizes[group_idx])
+            combo.append(sel_idx)
+
+        combo = tuple(combo)  # dictのキーにするのでtuple化
+        combination_to_filename[combo] = photo_filename
+
+    return combination_to_filename, group_sizes
+
+
+
+def main_offline_precompute():
+    """
+    1. Dataset(JSON)から (組み合わせ -> filename) を作る
+    2. CLIP モデルロード
+    3. filename -> embedding
+    4. ペアワイズ類似度
+    5. 保存
+    """
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_dir", type=str, default="lapwing/images")
+    parser.add_argument("--meta_file", type=str, default="lapwing/texts/BlendShapeGroupsMeta.json")
+    parser.add_argument("--teacher_data_file", type=str, default="lapwing/texts/BlendShapeData.json")
+    parser.add_argument("--clip_model_name", type=str, default="ViT-L-14")
+    parser.add_argument("--clip_pretrained", type=str, default="openai")
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--out_comb2fn", type=str, default="combination_to_filename.json")
+    parser.add_argument("--out_sims", type=str, default="pairwise_clip_sims.json")
+
+    args = parser.parse_args()
+
+    # 1. 組み合わせ->filename辞書の作成
+    combination_to_filename, group_sizes = build_combination_to_filename(
+        teacher_data_file=args.teacher_data_file,
+        meta_file=args.meta_file,
+    )
+
+
+    # 5. 保存 (JSON形式)
+    #   5.1 (combination -> filename)
+    #       group_sizes も保存しておくと後段のオンライン時に参照しやすい
+    #       tupleは文字列化する必要あり
+    comb2fn_dict = {
+        "group_sizes": group_sizes,
+        "mapping": {
+            ",".join(map(str, comb)): fn
+            for comb, fn in combination_to_filename.items()
+        }
+    }
+    with open(args.out_comb2fn, "w", encoding="utf-8") as f:
+        json.dump(comb2fn_dict, f, ensure_ascii=False, indent=2)
+
+    # データ出力
+    for combo, filename in combination_to_filename.items():
+        logger.info(f"Combination: {combo}, Filename: {filename}")
+
+
+if __name__ == "__main__":
+    main_offline_precompute()

combined_data/combination_to_filename.json

@@ -0,0 +1,80 @@
+{
+  "group_sizes": [
+    8,
+    9
+  ],
+  "mapping": {
+    "-1,-1": "000001.png",
+    "-1,0": "000002.png",
+    "-1,1": "000003.png",
+    "-1,2": "000004.png",
+    "-1,3": "000005.png",
+    "-1,4": "000006.png",
+    "-1,5": "000007.png",
+    "-1,6": "000008.png",
+    "-1,7": "000009.png",
+    "0,-1": "000010.png",
+    "0,0": "000011.png",
+    "0,1": "000012.png",
+    "0,2": "000013.png",
+    "0,3": "000014.png",
+    "0,4": "000015.png",
+    "0,5": "000016.png",
+    "0,6": "000017.png",
+    "0,7": "000018.png",
+    "1,-1": "000019.png",
+    "1,0": "000020.png",
+    "1,1": "000021.png",
+    "1,2": "000022.png",
+    "1,3": "000023.png",
+    "1,4": "000024.png",
+    "1,5": "000025.png",
+    "1,6": "000026.png",
+    "1,7": "000027.png",
+    "2,-1": "000028.png",
+    "2,0": "000029.png",
+    "2,1": "000030.png",
+    "2,2": "000031.png",
+    "2,3": "000032.png",
+    "2,4": "000033.png",
+    "2,5": "000034.png",
+    "2,6": "000035.png",
+    "2,7": "000036.png",
+    "3,-1": "000037.png",
+    "3,0": "000038.png",
+    "3,1": "000039.png",
+    "3,2": "000040.png",
+    "3,3": "000041.png",
+    "3,4": "000042.png",
+    "3,5": "000043.png",
+    "3,6": "000044.png",
+    "3,7": "000045.png",
+    "4,-1": "000046.png",
+    "4,0": "000047.png",
+    "4,1": "000048.png",
+    "4,2": "000049.png",
+    "4,3": "000050.png",
+    "4,4": "000051.png",
+    "4,5": "000052.png",
+    "4,6": "000053.png",
+    "4,7": "000054.png",
+    "5,-1": "000055.png",
+    "5,0": "000056.png",
+    "5,1": "000057.png",
+    "5,2": "000058.png",
+    "5,3": "000059.png",
+    "5,4": "000060.png",
+    "5,5": "000061.png",
+    "5,6": "000062.png",
+    "5,7": "000063.png",
+    "6,-1": "000064.png",
+    "6,0": "000065.png",
+    "6,1": "000066.png",
+    "6,2": "000067.png",
+    "6,3": "000068.png",
+    "6,4": "000069.png",
+    "6,5": "000070.png",
+    "6,6": "000071.png",
+    "6,7": "000072.png"
+  }
+}