# -*- coding: utf-8 -*- """End-to-End: 動画 → 多フレーム抽出 → 湯気除外 → スープ抽出 → 濃度推定 → 人間ラベルと比較。 以前作成した3モデルを連結する: 1) 湯気分類器 (EfficientNet-B0): 湯気が多いフレームを除外 2) RF-DETR セグメンテーション: 残ったフレームから soup_valid_area を抽出 3) 濃度分類器 (ResNet18, 10〜14): 抽出スープ画像から濃度ラベルを推定 フォールバック: - 全フレームが湯気多で除外された場合 → 最も湯気が薄い1枚を強制採用 - スープ抽出に失敗した場合 → フレーム中央クロップで代替 動画ごとに、採用フレームの濃度推定を信頼度で加重投票して最終ラベルを決め、 CSVの「スープ濃度」(人間ラベル) と突き合わせる。 """ from __future__ import annotations import argparse import csv import glob import json import subprocess import sys import tempfile from collections import defaultdict from pathlib import Path from typing import Any 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(checkpoint_path, device): import torch try: import timm model = timm.create_model("efficientnet_b0", pretrained=False, num_classes=2) except Exception: from torchvision import models import torch.nn as nn model = models.efficientnet_b0(weights=None) model.classifier[-1] = nn.Linear(model.classifier[-1].in_features, 2) ck = torch.load(checkpoint_path, map_location=device) model.load_state_dict(ck["model_state_dict"] if "model_state_dict" in ck else ck) model.to(device).eval() return model 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(video_path): r = subprocess.run(["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", str(video_path)], capture_output=True, text=True, timeout=30) try: return float(r.stdout.strip()) except ValueError: return 0.0 def extract_frame(video_path, seek): with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp: tp = Path(tmp.name) try: r = subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-ss", f"{seek:.3f}", "-i", str(video_path), "-frames:v", "1", "-q:v", "2", "-y", str(tp)], capture_output=True, text=True, timeout=60) if r.returncode != 0 or not tp.exists(): return None return cv2.imdecode(np.frombuffer(tp.read_bytes(), dtype=np.uint8), cv2.IMREAD_COLOR) finally: try: tp.unlink() except FileNotFoundError: pass def imwrite_u(path, image): path.parent.mkdir(parents=True, exist_ok=True) ok, enc = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), 92]) if ok: path.write_bytes(enc.tobytes()) def center_crop_bgr(image, ratio=0.6): h, w = image.shape[:2] cw, ch = int(w * ratio), int(h * ratio) x1, y1 = (w - cw) // 2, int(h * 0.30) return image[y1:min(h, y1 + ch), x1:x1 + cw] def resolve_video(video_dir, vid): for ext in VIDEO_EXTENSIONS: for c in (video_dir / f"{vid}{ext}", video_dir / f"{vid}{ext.upper()}"): if c.exists(): return c return None def load_human_labels(csv_glob): files = glob.glob(csv_glob) labels = {} if not files: return labels for row in csv.DictReader(open(files[0], encoding="utf-8-sig")): vid = (row.get("ID") or "").strip() dens = (row.get("スープ濃度") or "").strip() if vid: labels[vid] = dens return labels # ---- メイン ---------------------------------------------------------------- # 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=12) ap.add_argument("--steam-threshold", type=float, default=0.5) ap.add_argument("--crop-threshold", type=float, default=0.30) ap.add_argument("--label-csv", default=str(ROOT / "学習データ" / "*グリッド 10.csv")) ap.add_argument("--output-dir", default=str(ROOT / "local_batch_api" / "e2e_density")) 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_labels(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: vpath = resolve_video(video_dir, vid) if vpath is None: print(f"SKIP {vid}: not found"); continue dur = probe_duration(vpath) ratios = list(np.linspace(0.08, 0.92, args.frames)) print(f"\n=== {vid} ({vpath.name}) dur={dur:.1f}s ===") # 1) 全フレーム抽出 + 湯気判定 frames = [] # (idx, image, steam_prob, frame_path) for i, ratio in enumerate(ratios): img = extract_frame(vpath, max(0.0, dur * ratio)) if img is None: continue fp = out / "frames" / f"{vid}_f{i:02d}.jpg" imwrite_u(fp, img) frames.append([i, img, steam_prob(steam_model, steam_tf, img, device), fp]) if not frames: summary_rows.append(summary_row(vid, vpath, dur, args.frames, human, [], "no_frames")) continue # 2) 採用フレーム決定（フォールバック: 全除外なら最薄1枚） kept = [f for f in frames if f[2] < args.steam_threshold] fallback_steam = False if not kept: kept = [min(frames, key=lambda f: f[2])] fallback_steam = True # 3) スープ抽出 + 濃度推定 preds = [] # (label_id_int, confidence, source) for idx, img, sp, fp in kept: result = 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) status = int(result["extract_status"]) if status > 0: crop_path = result["soup_crop_path"]; src = "soup_crop" else: # フォールバック: 中央クロップ crop_path = str(crops_dir / f"{vid}_f{idx:02d}_centercrop.jpg") imwrite_u(Path(crop_path), center_crop_bgr(img)); src = "center_crop" d = predict_density(crop_path, str(DENSITY_CKPT)) lab = int(d["predicted_label"]); conf = float(d["confidence"]) preds.append((lab, conf, src)) frame_rows.append({"id": vid, "frame_index": idx, "steam_prob": round(sp, 4), "steam_fallback": int(fallback_steam), "crop_source": src, "crop_confidence": round(float(result["extract_confidence"]), 4), "density_pred": lab, "density_conf": round(conf, 4), "crop_path": crop_path}) print(f" f{idx:02d} steam={sp:.3f} crop={src} -> 濃度{lab} (conf {conf:.2f})") summary_rows.append(summary_row(vid, vpath, dur, args.frames, human, preds, "steam_fallback" if fallback_steam else "ok")) write_csv(out / "frame_density.csv", frame_rows) write_csv(out / "video_density_summary.csv", summary_rows) # 精度サマリ print("\n================ 動画ごと 濃度推定 vs 人間ラベル ================") diffs = [] exact = 0 n_eval = 0 for s in summary_rows: h = s["human_label"]; p = s["pred_label"] mark = "" if h not in ("", None) and p not in ("", None): n_eval += 1 d = abs(int(p) - int(h)); diffs.append(d) if d == 0: exact += 1; mark = "✓一致" elif d == 1: mark = "±1" else: mark = f"差{d}" print(f"{s['id']}: 人間={h} / 推定={p} (採用{s['frames_used']}枚 {s['note']}) {mark}") if n_eval: print(f"\n完全一致 {exact}/{n_eval} ({exact/n_eval:.0%}) / " f"±1以内 {sum(1 for d in diffs if d<=1)}/{n_eval} ({sum(1 for d in diffs if d<=1)/n_eval:.0%}) / " f"平均絶対誤差 MAE={np.mean(diffs):.2f}") print(f"\noutput -> {out}") return 0 def summary_row(vid, vpath, dur, frames_total, human, preds, note): pred_label = "" pred_conf = "" if preds: # 信頼度で加重投票 score = defaultdict(float) for lab, conf, _ in preds: score[lab] += conf best = max(score.items(), key=lambda kv: kv[1]) pred_label = best[0] pred_conf = round(best[1] / sum(score.values()), 3) return { "id": vid, "video": vpath.name, "duration_sec": round(dur, 1), "frames_total": frames_total, "frames_used": len(preds), "human_label": human.get(vid, ""), "pred_label": pred_label, "pred_vote_share": pred_conf, "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())