# -*- coding: utf-8 -*- """End-to-End v2: 選別基準を「スープが実際に見えているか(RF-DETR抽出信頼度)」に変更。 v1の課題: 「湯気の薄さ」で選ぶと、湯気が薄くてもスープ表面が見えないフレームを拾い、 center_cropフォールバックで濃度を過小評価していた。 v2の方針: 1) 各動画から多めに抽出 (既定20枚) 2) 全フレームにRF-DETRを掛け、抽出信頼度を測る (湯気が濃い=抽出失敗 で実質的な湯気ゲートになる) 3) 抽出信頼度 >= conf_keep のフレームのみ採用 → スープ表面が見えている瞬間だけ使う 4) 採用フレームに濃度推定 → (抽出信頼度 × 濃度信頼度) で加重投票 5) 採用ゼロの場合のみ、湯気最薄フレームのcenter_cropで保険推定 CSVの「スープ濃度」(人間ラベル) と突き合わせる。 """ from __future__ import annotations import argparse import csv import glob import subprocess import sys import tempfile from collections import defaultdict from pathlib import Path import cv2 import numpy as np ROOT = Path(__file__).resolve().parents[1] SUITE = Path.home() / "Documents" / "開発しているもの" / "制作中" / "AIラーメン開発" / "ramen-soup-ai-suite" CROP_SRC = SUITE / "03_soup-density-ai" / "src" CROP_CKPT = SUITE / "03_soup-density-ai" / "models" / "crop" / "soup_valid_area_rfdetr_seg_v2" / "checkpoint_best_ema.pth" DENSITY_CKPT = SUITE / "03_soup-density-ai" / "models" / "density" / "density_classifier_extracted_v1" / "model.pt" STEAM_CKPT = SUITE / "04_steam-classification" / "steam_detection_project" / "models" / "steam_classifier_best.pt" VIDEO_EXTENSIONS = [".mov", ".mp4", ".m4v", ".avi", ".mkv"] def build_steam_model(ckpt, device): import torch try: import timm m = timm.create_model("efficientnet_b0", pretrained=False, num_classes=2) except Exception: from torchvision import models import torch.nn as nn m = models.efficientnet_b0(weights=None) m.classifier[-1] = nn.Linear(m.classifier[-1].in_features, 2) c = torch.load(ckpt, map_location=device) m.load_state_dict(c["model_state_dict"] if "model_state_dict" in c else c) m.to(device).eval() return m def steam_transform(size=224): import albumentations as A from albumentations.pytorch import ToTensorV2 return A.Compose([A.Resize(height=size, width=size), A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2()]) def steam_prob(model, tf, image_bgr, device): import torch rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) t = tf(image=rgb)["image"].unsqueeze(0).to(device) with torch.no_grad(): return float(torch.softmax(model(t), dim=1)[0, 1].item()) def probe_duration(v): r = subprocess.run(["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", str(v)], capture_output=True, text=True, timeout=30) try: return float(r.stdout.strip()) except ValueError: return 0.0 def extract_frame(v, s): t = Path(tempfile.mktemp(suffix=".jpg")) subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-ss", f"{s:.3f}", "-i", str(v), "-frames:v", "1", "-q:v", "2", "-y", str(t)], capture_output=True, timeout=60) if t.exists(): img = cv2.imdecode(np.frombuffer(t.read_bytes(), np.uint8), cv2.IMREAD_COLOR) t.unlink() return img return None def imwrite_u(p, img): p.parent.mkdir(parents=True, exist_ok=True) ok, e = cv2.imencode(".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), 92]) if ok: p.write_bytes(e.tobytes()) def center_crop_bgr(img, ratio=0.6): h, w = img.shape[:2] cw, ch = int(w * ratio), int(h * ratio) x1, y1 = (w - cw) // 2, int(h * 0.30) return img[y1:min(h, y1 + ch), x1:x1 + cw] def resolve_video(d, vid): for ext in VIDEO_EXTENSIONS: for c in (d / f"{vid}{ext}", d / f"{vid}{ext.upper()}"): if c.exists(): return c return None def load_human(g): fs = glob.glob(g) out = {} if fs: for r in csv.DictReader(open(fs[0], encoding="utf-8-sig")): vid = (r.get("ID") or "").strip() if vid: out[vid] = (r.get("スープ濃度") or "").strip() return out def main(): ap = argparse.ArgumentParser() ap.add_argument("--video-dir", default=str(ROOT / "学習データ")) ap.add_argument("--ids", default="191-200") ap.add_argument("--frames", type=int, default=20) ap.add_argument("--conf-keep", type=float, default=0.50, help="この抽出信頼度以上のフレームを採用") ap.add_argument("--crop-threshold", type=float, default=0.10) ap.add_argument("--label-csv", default=str(ROOT / "学習データ" / "*グリッド 10.csv")) ap.add_argument("--output-dir", default=str(ROOT / "local_batch_api" / "e2e_density_v2")) args = ap.parse_args() import torch device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"device={device}") video_dir = Path(args.video_dir) out = Path(args.output_dir) crops_dir = out / "soup_crops" out.mkdir(parents=True, exist_ok=True) ids = parse_ids(args.ids) human = load_human(args.label_csv) print("loading models ...") steam_model = build_steam_model(STEAM_CKPT, device) steam_tf = steam_transform() sys.path.insert(0, str(CROP_SRC)) from crop.extract_soup_area import build_predictor, extract_one_image from density.predict_density import predict_density crop_predictor = build_predictor("small", str(CROP_CKPT), args.crop_threshold) frame_rows = [] summary_rows = [] for vid in ids: v = resolve_video(video_dir, vid) if v is None: print(f"SKIP {vid}"); continue dur = probe_duration(v) ratios = list(np.linspace(0.05, 0.95, args.frames)) print(f"\n=== {vid} ({v.name}) dur={dur:.1f}s ===") candidates = [] # (idx, crop_path, extract_conf, steam_prob) steam_all = [] # (idx, image, steam_prob, frame_path) best_extract = None # (idx, crop_path, extract_conf, steam_prob) 最も信頼度が高い抽出成功フレーム for i, r in enumerate(ratios): img = extract_frame(v, max(0.0, dur * r)) if img is None: continue fp = out / "frames" / f"{vid}_f{i:02d}.jpg" imwrite_u(fp, img) sp = steam_prob(steam_model, steam_tf, img, device) steam_all.append((i, img, sp, fp)) res = extract_one_image(image_path=fp, output_dir=crops_dir, checkpoint_path=str(CROP_CKPT), model_size="small", threshold=args.crop_threshold, low_confidence_threshold=args.crop_threshold, predictor=crop_predictor, save_alpha=False, preserve_filename=False, output_kind="crop", skip_fallback_output=True) ec = float(res["extract_confidence"]) if int(res["extract_status"]) > 0: if best_extract is None or ec > best_extract[2]: best_extract = (i, res["soup_crop_path"], ec, sp) if ec >= args.conf_keep: candidates.append((i, res["soup_crop_path"], ec, sp)) preds = [] note = "ok" if candidates: for idx, cp, ec, sp in candidates: d = predict_density(cp, str(DENSITY_CKPT)) lab, dc = int(d["predicted_label"]), float(d["confidence"]) preds.append((lab, ec * dc, "soup_crop")) frame_rows.append({"id": vid, "frame_index": idx, "steam_prob": round(sp, 4), "extract_conf": round(ec, 4), "crop_source": "soup_crop", "density_pred": lab, "density_conf": round(dc, 4)}) print(f" f{idx:02d} extract={ec:.2f} steam={sp:.2f} -> 濃度{lab} (conf {dc:.2f})") elif best_extract is not None: # フォールバックA: 0.5未満でも、最も信頼度が高い抽出成功フレームのスープ切り出しを使う note = "fallback_best_extract" idx, cp, ec, sp = best_extract d = predict_density(cp, str(DENSITY_CKPT)) lab, dc = int(d["predicted_label"]), float(d["confidence"]) preds.append((lab, dc, "best_extract")) frame_rows.append({"id": vid, "frame_index": idx, "steam_prob": round(sp, 4), "extract_conf": round(ec, 4), "crop_source": "best_extract", "density_pred": lab, "density_conf": round(dc, 4)}) print(f" [fallback:best_extract] f{idx:02d} extract={ec:.2f} -> 濃度{lab} (conf {dc:.2f})") else: # フォールバックB: 一度も抽出できなかった動画のみ中央クロップ note = "fallback_centercrop" idx, img, sp, fp = min(steam_all, key=lambda x: x[2]) if steam_all else (0, None, 1.0, None) if img is not None: cp = crops_dir / f"{vid}_f{idx:02d}_centercrop.jpg" imwrite_u(cp, center_crop_bgr(img)) d = predict_density(str(cp), str(DENSITY_CKPT)) lab, dc = int(d["predicted_label"]), float(d["confidence"]) preds.append((lab, dc, "center_crop")) frame_rows.append({"id": vid, "frame_index": idx, "steam_prob": round(sp, 4), "extract_conf": 0.0, "crop_source": "center_crop", "density_pred": lab, "density_conf": round(dc, 4)}) print(f" [fallback:center] f{idx:02d} steam={sp:.2f} -> 濃度{lab} (conf {dc:.2f})") summary_rows.append(summary_row(vid, v, dur, args.frames, human, preds, note, len(candidates))) write_csv(out / "frame_density.csv", frame_rows) write_csv(out / "video_density_summary.csv", summary_rows) print("\n================ v2: 濃度推定 vs 人間ラベル ================") diffs, exact, n = [], 0, 0 for s in summary_rows: h, p = s["human_label"], s["pred_label"] mark = "" if h not in ("", None) and p not in ("", None): n += 1 dd = abs(int(p) - int(h)); diffs.append(dd) mark = "✓一致" if dd == 0 else ("±1" if dd == 1 else f"差{dd}") if dd == 0: exact += 1 print(f"{s['id']}: 人間={h} / 推定={p} (採用{s['frames_used']}枚/{s['note']}) {mark}") if n: print(f"\n完全一致 {exact}/{n} ({exact/n:.0%}) / " f"±1以内 {sum(1 for d in diffs if d<=1)}/{n} ({sum(1 for d in diffs if d<=1)/n:.0%}) / " f"MAE={np.mean(diffs):.2f}") print(f"\noutput -> {out}") return 0 def summary_row(vid, v, dur, ft, human, preds, note, ncand): pl, share = "", "" if preds: sc = defaultdict(float) for lab, w, _ in preds: sc[lab] += w best = max(sc.items(), key=lambda kv: kv[1]) pl = best[0]; share = round(best[1] / sum(sc.values()), 3) return {"id": vid, "video": v.name, "duration_sec": round(dur, 1), "frames_total": ft, "candidates": ncand, "frames_used": len(preds), "human_label": human.get(vid, ""), "pred_label": pl, "vote_share": share, "frame_preds": ",".join(str(l) for l, _, _ in preds), "note": note} def parse_ids(spec): spec = spec.strip() if "-" in spec and "," not in spec: lo, hi = spec.split("-", 1) return [str(n) for n in range(int(lo), int(hi) + 1)] return [s.strip() for s in spec.split(",") if s.strip()] def write_csv(path, rows): if not rows: path.write_text("", encoding="utf-8-sig"); return fields = [] for r in rows: for k in r: if k not in fields: fields.append(k) with path.open("w", encoding="utf-8-sig", newline="") as fh: w = csv.DictWriter(fh, fieldnames=fields); w.writeheader(); w.writerows(rows) if __name__ == "__main__": raise SystemExit(main())