CoolFace
Apppublic

Hamzah-ALQadasi/Video_Authenticity.Anamoly_Detection.Temporal_Modeling

sourceHugging Faceapache-2.0updated 10mo agoView on Hugging Face
0likes
reorder_frames_algorithm.py381 linesDownload Raw Back to root
1#!/usr/bin/env python32"""3Frame order reconstruction algorithm using MSE and greedy path construction.4 5Reconstructs temporal frame order from shuffled videos using grayscale MSE matrix,6MST diameter endpoints, and double-ended greedy path building with local refinement.7 8Usage:9  # Process shuffled videos and CSVs from shuffled_artifacts10  python reorder_frames_algorithm.py --csv_dir ./shuffled_artifacts/shuffled_CSVs --videos_dir ./shuffled_artifacts/shuffled_videos --out_dir ./shuffled_artifacts/ordered_CSVs11 12Note: To generate reordered videos from predictions, use generate_ordered_videos_from_predictions.py13"""14 15import argparse16import os17import glob18 19import cv220import numpy as np21import pandas as pd22import torch23 24 25# =========================26# Config27# =========================28 29DEVICE = "cuda" if torch.cuda.is_available() else "cpu"30IMG_SIZE = 6431VIDEO_EXTS = (".avi", ".mp4", ".mov", ".mkv")32 33# =========================34# Pairwise MSE on GPU35# =========================36 37def compute_mse_matrix(frames: torch.Tensor) -> torch.Tensor:38    """39    frames: [N, 1, H, W] on DEVICE40    Returns:41        mse[i,j]: mean squared error between frame i and j42    """43    N = frames.shape[0]44    flat = frames.view(N, -1).float()           # [N, D]45 46    sq = (flat ** 2).sum(dim=1, keepdim=True)   # [N,1]47    dist2 = sq + sq.t() - 2.0 * (flat @ flat.t())48    dist2 = torch.clamp(dist2, min=0.0)49 50    D = flat.shape[1]51    mse = dist2 / D52    mse.fill_diagonal_(0.0)53    return mse54 55# =========================56# Utils57# =========================58 59def _mst_endpoints_via_diameter(mse: torch.Tensor):60    """61    Build an MST on the dense MSE matrix (edge weights = mse).62    Return (u, v) = endpoints of the MST diameter (longest weighted path).63    """64    N = mse.shape[0]65    if N <= 1:66        return (0, 0)67 68    device = mse.device69    used = torch.zeros(N, dtype=torch.bool, device=device)70    dist = torch.full((N,), float('inf'), device=device)71    parent = torch.full((N,), -1, dtype=torch.long, device=device)72 73    # start Prim from node 074    used[0] = True75    dist = mse[0].clone()76    dist[0] = float('inf')77 78    for _ in range(N - 1):79        masked = dist.clone()80        masked[used] = float('inf')81        j = int(torch.argmin(masked).item())82        used[j] = True83 84        # relax edges to unused nodes85        w = mse[j]86        update_mask = (~used) & (w < dist)87        dist[update_mask] = w[update_mask]88        parent[update_mask] = j89 90    # build adjacency list of the MST91    adj = [[] for _ in range(N)]92    for v in range(1, N):93        u = int(parent[v].item())94        if u >= 0:95            w = float(mse[u, v].item())96            adj[u].append((v, w))97            adj[v].append((u, w))98 99    def _farthest(src: int):100        # single-source longest distances on a tree via DFS 101        distv = [-1.0] * N102        distv[src] = 0.0103        stack = [src]104        while stack:105            x = stack.pop()106            for y, w in adj[x]:107                if distv[y] < 0.0:108                    distv[y] = distv[x] + w109                    stack.append(y)110        far = max(range(N), key=lambda k: distv[k])111        return far, distv[far]112 113    a, _ = _farthest(0)114    b, _ = _farthest(a)115    return a, b116 117def double_ended_greedy_from_pair(left: int, right: int, mse: torch.Tensor):118    """119    Maintain a path [left ... right]. At each step, attach the unused frame120    with minimal MSE to either end (choose the cheaper side).121    """122    N = mse.shape[0]123    used = torch.zeros(N, dtype=torch.bool, device=mse.device)124    used[left] = True125    used[right] = True126 127    path = [left, right]128    inf = float('inf')129 130    for _ in range(N - 2):131        # best to left132        candL = mse[:, left].clone()133        candL[used] = inf134        kL = int(torch.argmin(candL).item())135        dL = float(candL[kL])136 137        # best to right138        candR = mse[:, right].clone()139        candR[used] = inf140        kR = int(torch.argmin(candR).item())141        dR = float(candR[kR])142 143        if dL <= dR:144            path.insert(0, kL)145            used[kL] = True146            left = kL147        else:148            path.append(kR)149            used[kR] = True150            right = kR151 152    return path153 154 155def parse_shuffled_list(s: str):156    """157    Parse 'shuffled_frames_list' column.158    Example cell:159        "130,288,254,17,63,..."160    """161    return [int(x) for x in str(s).split(",") if x.strip() != ""]162 163 164 165def find_video_path(video_id: str, videos_dir: str) -> str:166    """167    Resolve the video path for a given video_id.168 169    Tries:170      - videos_dir / "<video_id>"171      - videos_dir / "<video_id>.avi"172      - videos_dir / "<video_id>.*" where extension in VIDEO_EXTS173    """174    # direct exact path (some CSVs store full filename)175    direct = os.path.join(videos_dir, video_id)176    if os.path.isfile(direct):177        return direct178 179    # try with .avi extension180    direct_avi = direct + ".avi"181    if os.path.isfile(direct_avi):182        return direct_avi183 184    # fallback: any file that starts with video_id185    pattern = os.path.join(videos_dir, f"{video_id}*")186    candidates = [187        p for p in glob.glob(pattern)188        if os.path.splitext(p)[1].lower() in VIDEO_EXTS189    ]190 191    if not candidates:192        raise FileNotFoundError(193            f"No video file found for video_id={video_id} in {videos_dir}"194        )195 196    # deterministic choice197    candidates.sort(key=lambda x: (len(os.path.basename(x)), x))198    return candidates[0]199 200 201# =========================202# Video loading (grayscale)203# =========================204 205def load_video_gray(video_path: str, expected_num_frames: int = None) -> torch.Tensor:206    """207    Load frames from a shuffled video as grayscale,208    resize to IMG_SIZE, and send to DEVICE.209 210    Returns:211        frames: [N, 1, H, W] float32 in [0,1] on DEVICE212    """213    if not os.path.isfile(video_path):214        raise FileNotFoundError(f"Video not found: {video_path}")215 216    cap = cv2.VideoCapture(video_path)217    if not cap.isOpened():218        raise IOError(f"Cannot open video: {video_path}")219 220    frames = []221    while True:222        ok, frame = cap.read()223        if not ok:224            break225        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)226        gray = cv2.resize(gray, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_AREA)227        frames.append(gray)228    cap.release()229 230    if len(frames) == 0:231        raise ValueError(f"No frames read from {video_path}")232 233    if expected_num_frames is not None and len(frames) != expected_num_frames:234        print(235            f"[WARN] {os.path.basename(video_path)}: "236            f"expected_num_frames={expected_num_frames}, read={len(frames)}"237        )238 239    arr = np.stack(frames, axis=0)          # [N, H, W]240    t = torch.from_numpy(arr).float()       # [N, H, W]241    t = t.unsqueeze(1) / 255.0              # [N, 1, H, W] in [0,1]242    return t.to(DEVICE)243 244 245# =========================246# Path construction247# =========================248 249def build_best_path(mse: torch.Tensor):250    """Build temporal path using MST diameter endpoints and double-ended greedy growth."""251    N = mse.shape[0]252    if N <= 2:253        return list(range(N))254 255    # smart seed via MST diameter256    a, b = _mst_endpoints_via_diameter(mse)257 258    # grow from both ends259    path = double_ended_greedy_from_pair(a, b, mse)260    261    return path262 263 264# =========================265# Per-video prediction266# =========================267 268def predict_order_for_video(video_id: str,269                            shuffled_order,270                            videos_dir: str):271    """272    Pipeline for a single video_id:273      - load shuffled video frames274      - compute MSE matrix275      - build best greedy path276      - refine path277      - map positions to original frame indices278    """279    shuffled_order = list(shuffled_order)280    expected_num_frames = len(shuffled_order)281 282    video_path = find_video_path(video_id, videos_dir)283    frames = load_video_gray(video_path, expected_num_frames=expected_num_frames)284    frames = frames[:, 0:1, :, :]  # use only Y channel for MSE285    N = frames.shape[0]286 287    if N != expected_num_frames:288        print(289            f"[WARN] {video_id}: csv_frames={expected_num_frames}, "290            f"video_frames={N}. Using min of both."291        )292        m = min(expected_num_frames, N)293        shuffled_order = shuffled_order[:m]294        frames = frames[:m]295        N = m296 297    if N <= 1:298        return [int(x) for x in shuffled_order]299    300    mse = compute_mse_matrix(frames)301    path = build_best_path(mse)302 303    predicted = [int(shuffled_order[idx]) for idx in path]304    return predicted305 306# =========================307# Process all CSVs308# =========================309 310def process_all_csvs(csv_dir: str, videos_dir: str, out_dir: str):311    """312    For each CSV in csv_dir:313      - read video_id, shuffled_frames_list314      - compute predicted order for each video315      - write a prediction CSV with same filename into out_dir316    """317    os.makedirs(out_dir, exist_ok=True)318 319    csv_paths = sorted(glob.glob(os.path.join(csv_dir, "*.csv")))320    if not csv_paths:321        raise FileNotFoundError(f"No CSV files found in {csv_dir}")322 323    for csv_path in csv_paths:324        df = pd.read_csv(csv_path)325        rows = []326 327        if "video_id" not in df.columns or "shuffled_frames_list" not in df.columns:328            raise ValueError(329                f"CSV {csv_path} must contain 'video_id' and 'shuffled_frames_list' columns."330            )331 332        for _, row in df.iterrows():333            video_id = str(row["video_id"]).strip()334            shuffled_order = parse_shuffled_list(row["shuffled_frames_list"])335            pred = predict_order_for_video(video_id, shuffled_order, videos_dir)336            pred_str = ",".join(str(x) for x in pred)337            rows.append({"video_id": video_id, "predicted_frames_list": pred_str})338 339        out_csv = os.path.join(out_dir, os.path.basename(csv_path))340        pd.DataFrame(rows).to_csv(out_csv, index=False)341        print(f"[OK] {os.path.basename(csv_path)} -> {os.path.basename(out_csv)}")342 343 344# =========================345# CLI346# =========================347 348def parse_args():349    parser = argparse.ArgumentParser(350        description="Reconstruct frame order from shuffled videos "351                    "using grayscale MSE and CSV metadata."352    )353    parser.add_argument(354        "--csv_dir",355        type=str,356        required=True,357        help="Directory with shuffled CSV files (e.g. shuffled_csvs).",358    )359    parser.add_argument(360        "--videos_dir",361        type=str,362        required=True,363        help="Directory with shuffled videos (e.g. UCF101_videos_shuffled).",364    )365    parser.add_argument(366        "--out_dir",367        type=str,368        default="./shuffled_artifacts/ordered_CSVs",369        help="Output directory for prediction CSVs.",370    )371    return parser.parse_args()372 373 374def main():375    args = parse_args()376    process_all_csvs(args.csv_dir, args.videos_dir, args.out_dir)377 378 379if __name__ == "__main__":380    main()381