"""배경/소품 추출 — 씬 원본(3000자) + 앞쪽 씬(2000자) + 해당 씬 샷 포함, 체이닝. 배경과 소품을 병렬로 실행 (각각 체이닝은 순차). Usage: cd backend .venv/bin/python experiments/entity_loc_prop_shots.py \ --project-id 84e9e90b-... --episode-id 78bc78d1-... \ --shots experiments/shot_results_gemini-pro.json \ --model gpt """ import argparse import json import sys import time from concurrent.futures import ThreadPoolExecutor from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from app.modules.llm.llm_client import call_structured from app.modules.pdf_parser import extract_text_from_pdf BUNDLE_TARGET = 3000 REF_MAX = 2000 SYSTEM = """시나리오 분석 전문가. 시나리오를 읽고 이미지 생성(T2I)에 필요한 시각적 요소를 리스팅한다. 카메라로 찍을 수 있는 시각적 요소만 리스팅. 리스팅에만 집중. 회상, 상상, 꿈 등 카메라에 찍히는 요소는 모두 포함.""" LOCATION_PROMPT = """위 씬들에서 등장하는 배경(장소)을 모두 나열하세요. ## 리스팅 항목 - name: 배경 이름 - shot_count: 해당 배경이 등장하는 샷 수 ## 포함 기준 - 카메라에 찍히는 모든 장소/공간 - 실내/실외 모두 포함 - 같은 장소라도 시간대가 다르면 하나로 통합 ## 제외 기준 - 언급만 되고 시각적으로 보이지 않는 장소는 제외""" PROP_PROMPT = """위 씬들에서 등장하는 소품(물건)을 모두 나열하세요. ## 리스팅 항목 - name: 소품 이름 - shot_count: 해당 소품이 등장하는 샷 수 ## 포함 기준 - 3샷 이상 반복 등장 - 캐릭터가 착용/사용/탑승하는 것만 - 부속 소품은 본체에 통합 (아머+헬멧→아머) ## 제외 기준 - 일반 물건, 배경 시설, UI 화면, 의류, 상황 묘사 제외 - 에피소드당 5~10개 목표""" LOCATION_SCHEMA = { "type": "object", "properties": { "locations": { "type": "array", "items": { "type": "object", "properties": { "name": {"type": "string"}, "shot_count": {"type": "integer"}, }, "required": ["name", "shot_count"], "additionalProperties": False, }, }, }, "required": ["locations"], "additionalProperties": False, } PROP_SCHEMA = { "type": "object", "properties": { "props": { "type": "array", "items": { "type": "object", "properties": { "name": {"type": "string"}, "shot_count": {"type": "integer"}, }, "required": ["name", "shot_count"], "additionalProperties": False, }, }, }, "required": ["props"], "additionalProperties": False, } def load_scenes(project_dir: Path, episode_id: str): cp_path = (project_dir / "checkpoints" / "episodes" / episode_id / "scene_segmentation" / "manifest.json") cp = json.loads(cp_path.read_text()) segments = cp["data"]["segments"] pdf_dir = project_dir / "assets" / "screenplays" pdfs = list(pdf_dir.glob("*.pdf")) fulltext, _ = extract_text_from_pdf(pdfs[0]) scene_texts = [] for seg in segments: start = seg["start_char"] end = seg["end_char"] scene_texts.append({ "idx": seg["scene_index"], "heading": seg.get("heading", ""), "text": fulltext[start:end], "length": end - start, }) return scene_texts def load_shots(shots_path: str): shots_data = json.loads(Path(shots_path).read_text()) shots_by_scene = {} for s in shots_data: shots_by_scene[s["scene_index"]] = s.get("shots", []) return shots_by_scene def _format_shots(shots): if not shots: return "" lines = [] for sh in shots: chars = ", ".join(sh.get("characters", [])) lines.append(f" Shot {sh['shot_index']}: {sh['description'][:120]} [{chars}]") return "\n".join(lines) def extract_chained(scene_texts, shots_by_scene, model, entity_type): """체이닝 방식으로 배경 또는 소품 추출.""" if entity_type == "location": type_prompt = LOCATION_PROMPT schema = LOCATION_SCHEMA key = "locations" else: type_prompt = PROP_PROMPT schema = PROP_SCHEMA key = "props" step_name = f"entity_shot_{entity_type}" all_entities = [] existing_names = set() i = 0 call_count = 0 total_start = time.time() while i < len(scene_texts): # 번들 구성 (BUNDLE_TARGET) bundle = [] bundle_len = 0 while i < len(scene_texts) and bundle_len + scene_texts[i]["length"] <= BUNDLE_TARGET: bundle.append(scene_texts[i]) bundle_len += scene_texts[i]["length"] i += 1 if not bundle and i < len(scene_texts): bundle.append(scene_texts[i]) bundle_len = scene_texts[i]["length"] i += 1 # 앞쪽 참조 (REF_MAX) ref_parts = [] ref_len = 0 j = bundle[0]["idx"] - 2 while j >= 0 and j < len(scene_texts) and ref_len + scene_texts[j]["length"] <= REF_MAX: ref_parts.insert(0, scene_texts[j]["text"]) ref_len += scene_texts[j]["length"] j -= 1 # 프롬프트 구성 user_parts = [] # 체이닝 if all_entities: user_parts.append( "지금까지 찾은 요소:\n" + "\n".join(f"- {e['name']} (shot_count: {e['shot_count']})" for e in all_entities) + "\n\n위 목록에 없는 새로운 요소만 추가하세요.\n" ) # 앞쪽 씬 (참조만, 샷 없음) if ref_parts: user_parts.append(f"[앞쪽 씬 — 참조만]\n{''.join(ref_parts)}\n") # 분석 대상 씬 + 샷 for s in bundle: shots = shots_by_scene.get(s["idx"], []) shot_text = _format_shots(shots) scene_block = f"--- Scene {s['idx']}: {s['heading']} ---\n{s['text']}" if shot_text: scene_block += f"\n[샷]\n{shot_text}" user_parts.append(scene_block) user_parts.append(type_prompt) user_prompt = "\n\n".join(user_parts) call_count += 1 call_start = time.time() scene_range = f"{bundle[0]['idx']}~{bundle[-1]['idx']}" try: result = call_structured( step=step_name, system_prompt=SYSTEM, user_prompt=user_prompt, response_schema=schema, project_config={"default_model": model}, schema_name=f"{step_name}_{call_count}", ) elapsed = time.time() - call_start new_items = result.get(key, []) added = 0 for e in new_items: name = e.get("name", "") if not name: continue if name in existing_names: for existing in all_entities: if existing["name"] == name: new_sc = e.get("shot_count", 0) if new_sc > 0: existing["shot_count"] += new_sc break else: all_entities.append(e) existing_names.add(name) added += 1 print(f" [{entity_type}] call {call_count} (scenes {scene_range}, {elapsed:.1f}s): " f"+{added} (total {len(all_entities)})", flush=True) except Exception as exc: elapsed = time.time() - call_start print(f" [{entity_type}] call {call_count} (scenes {scene_range}) FAILED ({elapsed:.1f}s): {exc}", flush=True) total_elapsed = time.time() - total_start all_entities.sort(key=lambda x: x.get("shot_count", 0), reverse=True) print(f"\n [{entity_type}] Total: {len(all_entities)} items, {call_count} calls, {total_elapsed:.1f}s", flush=True) return all_entities def main(): parser = argparse.ArgumentParser() parser.add_argument("--project-id", required=True) parser.add_argument("--episode-id", required=True) parser.add_argument("--shots", required=True) parser.add_argument("--model", default="gpt") args = parser.parse_args() projects_root = Path(__file__).resolve().parent.parent.parent / "projects" project_dir = projects_root / args.project_id print("Loading scenes...") scene_texts = load_scenes(project_dir, args.episode_id) print(f" {len(scene_texts)} scenes") print(f"Loading shots from {args.shots}...") shots_by_scene = load_shots(args.shots) print(f" {sum(len(v) for v in shots_by_scene.values())} shots") print(f"\nExtracting with model={args.model} (location + prop 병렬)...\n") with ThreadPoolExecutor(max_workers=2) as pool: loc_future = pool.submit(extract_chained, scene_texts, shots_by_scene, args.model, "location") prop_future = pool.submit(extract_chained, scene_texts, shots_by_scene, args.model, "prop") locations = loc_future.result() props = prop_future.result() # 저장 for name, data in [("location", locations), ("prop", props)]: out = f"experiments/entity_shot_{name}_{args.model}.json" Path(out).write_text(json.dumps(data, ensure_ascii=False, indent=2)) print(f"\nSaved {name} to {out}") print(f"=== {name} ({len(data)} items) ===") for e in data: print(f" {e['name']:25s} | shots: {e.get('shot_count', 0)}") if __name__ == "__main__": main()