#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从 nyu_depth_v2_labeled.mat 导出官方 Eigen test split 的 654 张 RGB 图片。

划分来源：BTS 仓库 utils/splits.mat 的 testNdxs（1-based，共 654 个），
与 train_test_inputs/nyudepthv2_test_files_with_gt.txt 的 654 行逐场景帧数一致，
已核对无重叠、并集覆盖 1..1449。

只导出 RGB，GT 不落盘：评测脚本按 split.json 里的 mat_index 直接从 mat 读取，
避免重复存储。

用法：
    # 导出 654 张到 datasets/nyu654_test/
    ../../venv/bin/python src/tool/export_nyu654_test.py \
        --mat ../../datasets/nyu_depth_v2_labeled.mat \
        --splits datasets/nyu654_test/splits.mat \
        --out-dir datasets/nyu654_test

输出：
    <out-dir>/images/{序号:04d}_{scene}.png
    <out-dir>/split.json    [{mat_index, scene, file}, ...]
"""

import argparse
import json
from pathlib import Path

import cv2
import h5py
import numpy as np
import scipy.io as sio
from tqdm import tqdm


def decode_cell(f, ref):
    return ''.join(chr(c) for c in np.asarray(f[ref]).ravel())


def main():
    parser = argparse.ArgumentParser(description='导出 NYU Depth V2 Eigen test split 654 张')
    parser.add_argument('--mat', required=True, help='nyu_depth_v2_labeled.mat')
    parser.add_argument('--splits', required=True, help='含 testNdxs 的 splits.mat')
    parser.add_argument('--out-dir', required=True, help='输出目录')
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    image_dir = out_dir / 'images'
    image_dir.mkdir(parents=True, exist_ok=True)

    ndxs = sio.loadmat(args.splits)['testNdxs'].ravel().astype(int)
    if len(ndxs) != 654:
        raise ValueError(f'testNdxs 应为 654 个，实际 {len(ndxs)}')

    records = []
    with h5py.File(args.mat, 'r') as f:
        for seq, mat_index in enumerate(tqdm(ndxs, desc='export', ncols=80)):
            i = mat_index - 1  # matlab 1-based
            scene = decode_cell(f, f['scenes'][0, i])
            rgb = np.asarray(f['images'][i]).transpose(1, 2, 0)  # CHW -> HWC
            name = f'{seq:04d}_{scene}.png'
            if not cv2.imwrite(str(image_dir / name), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)):
                raise RuntimeError(f'写入失败: {image_dir / name}')
            records.append({'mat_index': int(i), 'scene': scene, 'file': name})

    with open(out_dir / 'split.json', 'w') as fp:
        json.dump(records, fp, indent=2)
    print(f'已导出 {len(records)} 张到 {image_dir}，划分文件 {out_dir / "split.json"}')


if __name__ == '__main__':
    main()