Add files using upload-large-folder tool

FE2E/.gitignore

@@ -0,0 +1,35 @@
+.DS_store
+.idea
+*/.DS_store
+__pycache__
+*/__pycache__/
+*/.ipynb_checkpoints/
+
+#特殊文件
+*.report
+*.png
+*-xdl
+
+#几个不改动或无意义改动文件
+Qwen/
+pretrain/
+splits/
+log_err/
+assets/
+!assets/
+!assets/demo.png
+!assets/pipeline.png
+vkitti/
+flux/
+compare_result/
+
+#推理用数据集
+infer/data_split/
+infer/nyu_results_test/
+infer/nyudepth/
+infer/eval_results/
+infer/diode/
+infer/eth3d/
+infer/kitti/
+infer/scannet/
+infer/dsine_eval/

FE2E/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 AMAP-ML
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

FE2E/README.md

@@ -0,0 +1,159 @@
+# FE2E: From Editor to Dense Geometry Estimator
+
+[![Page](https://img.shields.io/badge/Project-Website-pink?logo=googlechrome&logoColor=white)](https://amap-ml.github.io/FE2E/)
+[![Paper](https://img.shields.io/badge/arXiv-2509.04338-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2509.04338)
+[![GitHub](https://img.shields.io/github/stars/AMAP-ML/FE2E?style=social)](https://github.com/AMAP-ML/FE2E)
+[![HuggingFace](https://img.shields.io/badge/🤗%20HuggingFace-Model-yellow)](https://huggingface.co/exander/FE2E)
+[![Video](https://img.shields.io/badge/BiliBili-Video-00A1D6)](https://www.bilibili.com/video/BV1zYXdBXE2x)
+[![Video](https://img.shields.io/badge/YouTube-Video-red)](https://youtu.be/fyXwwH_-o5w)
+
+[Jiyuan Wang](https://wangjiyuan9.github.io/)<sup>1,2</sup>,
+[Chunyu Lin](https://scholar.google.com/citations?hl=zh-CN&user=t8xkhscAAAAJ)<sup>1&#9993;</sup>,
+[Lei Sun](https://scholar.google.com/citations?user=your-id)<sup>2&#10013;</sup>,
+[Rongying Liu](https://scholar.google.com/citations?user=your-id)<sup>1</sup>,
+[Mingxing Li](https://scholar.google.com/citations?user=-pfkprkAAAAJ&hl=zh-CN&oi=ao)<sup>2</sup>,
+[Lang Nie](https://scholar.google.com/citations?hl=zh-CN&user=vo__egkAAAAJ)<sup>3</sup>,
+[Kang Liao](https://kangliao929.github.io/)<sup>4</sup>,
+[Xiangxiang Chu](https://cxxgtxy.github.io/)<sup>2</sup>,
+[Yao Zhao](https://faculty.bjtu.edu.cn/5900/)<sup>1</sup>
+
+<span class="author-block"><sup>1</sup>Beijing Jiaotong University</span>  
+<span class="author-block"><sup>2</sup>Alibaba Group</span>  
+<span class="author-block"><sup>3</sup>Chongqing University of Posts and Telecommunications</span>  
+<span class="author-block"><sup>4</sup>Nanyang Technological University</span>  
+<span class="author-block"><sup>&#9993;</sup>Corresponding author. <sup>&#10013;</sup>Project leader.</span>
+
+![teaser](assets/demo.png)
+
+We present **FE2E**, a DiT-based foundation model for monocular dense geometry prediction. FE2E adapts an advanced image editing model to dense geometry tasks and achieves strong zero-shot performance on both monocular depth and normal estimation.
+
+![pipeline](assets/pipeline.png)
+
+## 📢 News
+- **[2026-03-17]**: Code and Checkpoint are available now!
+- **[2026-02-21]**: FE2E was accepted by CVPR 2026!!! 🎉🎉🎉
+- **[2025-09-05]**: Paper released on [arXiv](https://arxiv.org/abs/2509.04338).
+
+---
+
+## 🛠️ Setup
+
+This codebase is prepared as an inference/evaluation release.
+
+```bash
+pip install -r requirements.txt
+```
+
+Recommended local layout:
+
+```text
+FE2E/
+├── pretrain/
+│   ├── step1x-edit-i1258.safetensors
+│   ├── step1x-edit-v1p1-official.safetensors
+│   └── vae.safetensors
+├── lora/
+│   └── LDRN.safetensors
+├── infer/
+│   ├── eth3d/
+│   │   └── eth3d.tar
+│   └── dsine_eval/
+│       ├── nyuv2/
+│       └── scannet/
+└── logs/
+```
+
+---
+
+## 🔥 Training
+
+```text
+[ ] Training code will be released later.
+```
+
+---
+
+## 🕹️ Inference
+
+### 1. Prepare Model Weights
+
+1.  Download the base weights, which from the official [Step1X-Edit](https://github.com/stepfun-ai/Step1X-Edit) release.
+2.  Download FE2E LoRA [checkpoint](https://huggingface.co/exander/FE2E/blob/main/LDRN.safetensors)
+
+
+### 2. Prepare Benchmark Datasets
+
+- Depth benchmarks follow the external evaluation data convention from [Marigold](https://github.com/prs-eth/Marigold).
+- Normal benchmarks follow the external evaluation data convention from [DSINE](https://github.com/baegwangbin/DSINE).
+
+
+Supported depth benchmarks:
+- `nyu_v2`,`kitti`,`eth3d`,`diode`,`scannet`
+
+Supported normal benchmarks:
+- `nyuv2`,`scannet`,`ibims`,`sintel`
+
+
+### 3. Run Evaluation
+
+`[dataset] normal`:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+MASTER_PORT=21258 \
+PYTHONUNBUFFERED=1 \
+python -u evaluation.py \
+  --model_path ./pretrain \
+  --eval_data_root ./infer \
+  --output_dir ./infer/eval_verify_scannet_normal_8gpu \
+  --num_gpus 8 \
+  --num_samples -1 \
+  --lora ./lora/LDRN.safetensors \
+  --single_denoise \
+  --prompt_type empty \
+  --norm_type ln \
+  --task_name normal \
+  --normal_eval_datasets [dataset]
+```
+
+`[dataset] depth`:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+MASTER_PORT=21257 \
+PYTHONUNBUFFERED=1 \
+python -u evaluation.py \
+  --model_path ./pretrain \
+  --eval_data_root ./infer \
+  --output_dir ./infer/eval_verify_eth3d_8gpu \
+  --num_gpus 8 \
+  --num_samples -1 \
+  --lora ./lora/LDRN.safetensors \
+  --single_denoise \
+  --prompt_type empty \
+  --norm_type ln \
+  --task_name depth \
+  --depth_eval_datasets [dataset]
+```
+
+
+### 4. Reference Logs
+If you want to known the successful status, this repo includes run logs in `logs/`:
+- `logs/verify_scannet_normal_8gpu_20260317_171345.log`
+- `logs/verify_eth3d_8gpu_20260317_172004.log`
+
+
+---
+
+## 🎓 Citation
+
+If you find our work useful, please cite:
+
+```bibtex
+@article{wang2025editor,
+  title={From Editor to Dense Geometry Estimator},
+  author={Wang, JiYuan and Lin, Chunyu and Sun, Lei and Liu, Rongying and Nie, Lang and Li, Mingxing and Liao, Kang and Chu, Xiangxiang and Zhao, Yao},
+  journal={arXiv preprint arXiv:2509.04338},
+  year={2025}
+}
+```

FE2E/evaluation.py

@@ -0,0 +1,553 @@
+import argparse
+import os
+import re
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from infer.seed_all import seed_all
+
+# 设置环境变量消除tokenizers警告
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+os.environ['NCCL_DEBUG'] = 'WARN'
+# 消除torchvision警告
+os.environ['TORCHVISION_DISABLE_DEPRECATED_WARNING'] = '1'
+
+REPO_ROOT = os.path.dirname(os.path.abspath(__file__))
+DEFAULT_QWEN_DIR = os.path.join(REPO_ROOT, "Qwen")
+DEFAULT_DEPTH_DATASET_CONFIGS = {
+    "nyu_v2": "configs/data_nyu_test.yaml",
+    "kitti": "configs/data_kitti_eigen_test.yaml",
+    "eth3d": "configs/data_eth3d.yaml",
+    "diode": "configs/data_diode_all.yaml",
+    "scannet": "configs/data_scannet_val.yaml",
+}
+DEFAULT_NORMAL_DATASETS = {
+    "nyuv2": "test",
+    "scannet": "test",
+    "ibims": "ibims",
+    "sintel": "sintel",
+    "oasis": "val",
+    "hypersim": "hypersim",
+}
+
+
+def resolve_eval_data_root(args, *required_markers):
+    """Resolve the evaluation data root without depending on the launch cwd."""
+    candidates = []
+    if getattr(args, "eval_data_root", None):
+        candidates.append(os.path.abspath(args.eval_data_root))
+
+    candidates.extend(
+        [
+            os.path.join(REPO_ROOT, "infer"),
+            os.path.join(REPO_ROOT, "data"),
+            os.path.join(os.path.dirname(REPO_ROOT), "data"),
+        ]
+    )
+
+    for candidate in candidates:
+        if all(os.path.exists(os.path.join(candidate, marker)) for marker in required_markers):
+            return candidate
+
+    checked = ", ".join(
+        os.path.join(candidate, marker) for candidate in candidates for marker in required_markers
+    )
+    raise FileNotFoundError(f"未找到评测数据根目录，已检查: {checked}")
+
+
+def parse_depth_eval_datasets(raw_value):
+    requested = [item.strip() for item in raw_value.split(",") if item.strip()]
+    if requested == ["all"]:
+        requested = list(DEFAULT_DEPTH_DATASET_CONFIGS.keys())
+    invalid = [item for item in requested if item not in DEFAULT_DEPTH_DATASET_CONFIGS]
+    if invalid:
+        raise ValueError(f"不支持的 depth 数据集: {invalid}，可选值: {sorted(DEFAULT_DEPTH_DATASET_CONFIGS)}")
+    return {name: DEFAULT_DEPTH_DATASET_CONFIGS[name] for name in requested}
+
+
+def parse_normal_eval_datasets(raw_value):
+    requested = [item.strip() for item in raw_value.split(",") if item.strip()]
+    if requested == ["all"]:
+        requested = list(DEFAULT_NORMAL_DATASETS.keys())
+    invalid = [item for item in requested if item not in DEFAULT_NORMAL_DATASETS]
+    if invalid:
+        raise ValueError(f"不支持的 normal 数据集: {invalid}，可选值: {sorted(DEFAULT_NORMAL_DATASETS)}")
+    return [(name, DEFAULT_NORMAL_DATASETS[name]) for name in requested]
+
+
+def collect_and_merge_dual_cfg_results(rank, world_size, gathered_metrics_Lpred, gathered_times):
+    """
+    收集并合并双CFG配置的评估结果
+    
+    Args:
+        rank: 当前进程的rank
+        world_size: 总进程数
+        gathered_metrics_Lpred
+        gathered_times: 处理时间收集结果
+    
+    Returns:
+        tuple: (all_metrics_Lpred, dataset_times)
+    """
+    if rank != 0:
+        return None, None
+
+    # 合并处理时间
+    dataset_times = []
+    for times_list in gathered_times:
+        dataset_times.extend(times_list)
+
+    #先处理L的
+    all_metrics_L = {}
+    valid_metrics_L = [m for m in gathered_metrics_Lpred if m]
+    if valid_metrics_L:
+        for key in valid_metrics_L[0].keys():
+            values = [m[key] for m in valid_metrics_L if key in m]
+            if values:
+                all_metrics_L[key] = np.mean(values)
+
+    return all_metrics_L, dataset_times
+
+
+def format_dual_cfg_results_table(dataset_name, model_identifier, all_metrics_L, dataset_times):
+    """
+    格式化双CFG配置的结果表格
+    
+    Args:
+        dataset_name: 数据集名称
+        model_identifier: 模型标识符
+        all_metrics_L: CFG=1的评估指标
+        dataset_times: 处理时间列表
+    
+    Returns:
+        str: 格式化的结果字符串
+    """
+    eval_metrics_order = ["abs_relative_difference", "squared_relative_difference", "rmse_linear", "rmse_log", "delta1_acc", "delta2_acc", "delta3_acc"]
+
+    # 获取CFG=1的指标值
+    mean_errors_L = [all_metrics_L.get(metric, 0.0) for metric in eval_metrics_order]
+
+    # 构建表格
+    metrics_header = ["Dataset", "Model", "CFG", "abs_rel", "sq_rel", "rmse", "rmse_log", "a1", "a2", "a3"]
+
+    # CFG=1的结果行
+    values_data_L = [dataset_name, model_identifier, "CFG=1"] + [f"{v:.4f}" for v in mean_errors_L]
+
+    header_line = "| " + " | ".join(metrics_header) + " |"
+    separator_line = "| " + " | ".join(["---"] * len(metrics_header)) + " |"
+    values_line_L = "| " + " | ".join(values_data_L) + " |"
+
+    # 生成输出字符串
+    result_str = f"\n数据集 {dataset_name} 评估完成!\n"
+    result_str += "-" * 100 + "\n"
+    result_str += header_line + "\n"
+    result_str += separator_line + "\n"
+    result_str += values_line_L + "\n"
+
+    # 添加统计信息
+    sample_count = len(dataset_times)
+    result_str += f"样本数: {sample_count}\n"
+    if dataset_times:
+        result_str += f"平均处理时间: {np.mean(dataset_times):.2f}秒/图像\n"
+    result_str += "-" * 100 + "\n"
+
+    return result_str
+
+
+def save_dual_cfg_results_summary(output_dir, all_dataset_results, model_identifier):
+    """
+    保存所有数据集的双CFG评估结果摘要
+    
+    Args:
+        output_dir: 输出目录
+        all_dataset_results: 所有数据集的结果字典
+    """
+    summary_file = os.path.join(output_dir, f"{model_identifier}.txt")
+
+    with open(summary_file, 'w', encoding='utf-8') as f:
+        f.write("=" * 120 + "\n")
+        f.write("双CFG配置深度评估结果汇总\n")
+        f.write("=" * 120 + "\n\n")
+
+        for dataset_name, result_data in all_dataset_results.items():
+            f.write(result_data['formatted_output'])
+            f.write(f"结果保存至: {result_data['eval_dir']}\n\n")
+
+    print(f"双CFG评估结果摘要已保存至: {summary_file}")
+
+
+def parse_args():
+    '''Set the Args'''
+    parser = argparse.ArgumentParser(description="Run Step...")
+    parser.add_argument('--model_path', type=str, default='./pretrain', help='模型路径')
+    parser.add_argument('--qwen2vl_model_path', type=str, default=DEFAULT_QWEN_DIR, help='Qwen2.5-VL 模型目录')
+    parser.add_argument('--eval_data_root', type=str, default=None, help='评测数据根目录，默认自动在仓库相对路径下查找')
+    parser.add_argument("--seed", type=int, default=1234, help="随机种子")
+    parser.add_argument("--output_dir", type=str, default="./infer/eval_results", help="Output directory.")
+    parser.add_argument('--num_steps', type=int, default=28, help='扩散步数')
+    parser.add_argument('--num_samples', type=int, default=-1, help='生成样本数')
+    parser.add_argument('--cfg_guidance', type=float, default=6.0, help='CFG引导强度')
+    parser.add_argument('--size_level', type=int, default=768, help='输入图像大小')
+    parser.add_argument('--num_gpus', type=int, default=torch.cuda.device_count(), help='使用的GPU数量')
+    parser.add_argument('--save_viz', action='store_true', help='保存可视化结果')
+    parser.add_argument('--offload', action='store_true', help='使用CPU卸载以节省GPU内存')
+    parser.add_argument('--quantized', action='store_true', help='使用量化模型')
+    parser.add_argument('--lora', type=str, help='LoRA模型路径')
+    parser.add_argument('--single_denoise', action='store_true', default=False, help='单步推理')
+    parser.add_argument('--old_prompt', action='store_true', default=False, help='使用旧版提示')
+    parser.add_argument('--prompt_type', type=str, default='query', help='提示类型')
+    parser.add_argument('--prompt', type=str, default='Describe the 3D structure and layout of the scene in the image. Predict the depth of this image.', help='提示')
+    parser.add_argument('--norm_type', type=str, default='depth', help='预测结果的归一化方式，目前有 depth、disp、ln')
+    parser.add_argument('--task_name', type=str, default='depth', help='任务名称，支持 depth 或 normal')
+    parser.add_argument('--depth_eval_datasets', type=str, default='eth3d', help='逗号分隔的 depth 评测数据集')
+    parser.add_argument('--normal_eval_datasets', type=str, default='nyuv2,scannet', help='逗号分隔的 normal 评测数据集')
+    parser.add_argument('--debug', action='store_true', default=False, help='调试模式')
+    args = parser.parse_args()
+    if args.single_denoise:
+        args.num_steps = 1
+    return args
+
+
+def extract_model_identifier(lora_path):
+    """
+    从lora路径中提取模型标识符
+    支持多种路径格式：
+    - ./log_err/dis-hvsge-log/ckpt.safetensors -> dis-hvsge-log
+    - /path/to/folder/ckpt-123 -> folder-epoch123
+    - /path/to/model.safetensors -> model
+    """
+    if not lora_path or not os.path.exists(lora_path):
+        return "DefaultModel"
+    
+    # 规范化路径
+    lora_path = os.path.normpath(lora_path)
+    
+    # 方法1: 匹配 /folder/ckpt-数字 格式
+    match = re.search(r'/([^/]+)/ckpt-(\d+)', lora_path)
+    if match:
+        folder_name = match.group(1)
+        epoch_num = int(match.group(2))
+        return f"{folder_name}-epoch{epoch_num}"
+    
+    # 方法2: 匹配 /folder/ckpt.safetensors 格式
+    match = re.search(r'/([^/]+)/ckpt\.safetensors$', lora_path)
+    if match:
+        return match.group(1)
+    
+    # 方法3: 匹配 ./folder/ckpt.safetensors 格式
+    match = re.search(r'[./]*([^/]+)/ckpt\.safetensors$', lora_path)
+    if match:
+        return match.group(1)
+    
+    # 方法4: 从文件名中提取（兜底方案）
+    filename = os.path.basename(lora_path)
+    return filename.split('.')[0] if '.' in filename else filename
+
+
+def setup(rank, world_size):
+    """初始化分布式环境"""
+    os.environ.setdefault('MASTER_ADDR', 'localhost')
+    os.environ.setdefault('MASTER_PORT', '21256')
+    dist.init_process_group("nccl", rank=rank, world_size=world_size)
+
+
+def cleanup():
+    """清理分布式环境"""
+    dist.destroy_process_group()
+
+
+def main_worker(rank, world_size, args, dataset_configs):
+    """每个进程的主函数"""
+    from infer.inference import ImageGenerator
+    from infer.inner_evaluation import evaluation_depth_custom_parallel
+
+    setup(rank, world_size)
+    torch.cuda.set_device(rank)
+    device = torch.device(f"cuda:{rank}")
+
+    if rank == 0:
+        print(f"[main_worker] 开始加载 pipeline, device={device}, datasets={list(dataset_configs.keys())}", flush=True)
+
+    pipeline = ImageGenerator(
+        ae_path=os.path.join(args.model_path, 'vae.safetensors'),
+        dit_path=os.path.join(args.model_path, "step1x-edit-i1258-FP8.safetensors" if args.quantized else "step1x-edit-i1258.safetensors"),
+        qwen2vl_model_path=args.qwen2vl_model_path,
+        max_length=640,
+        quantized=args.quantized,
+        offload=args.offload,
+        lora=args.lora,
+        device=str(device),
+        args=args,
+    )
+
+    if rank == 0:
+        print(f"Successfully loading pipeline from {args.model_path}.", flush=True)
+
+    test_data_dir = resolve_eval_data_root(args, "configs")
+
+    # 使用新的模型标识符提取函数
+    model_identifier = extract_model_identifier(args.lora)
+
+    if rank == 0:
+        print(f"模型标识符: {model_identifier}", flush=True)
+
+    all_dataset_results = {}
+    aligment_map = {"depth": "least_square", "disp": "least_square_disparity", "ln": "log_space"}
+    for dataset_name, config_path in dataset_configs.items():
+        # 修改输出目录结构：在数据集名称外添加模型名称层级
+        eval_dir = os.path.join(args.output_dir, model_identifier, dataset_name)
+        test_dataset_config = os.path.join(test_data_dir, config_path)
+        alignment_type = aligment_map[args.norm_type]
+
+        if rank == 0:
+            print(f"\n开始评估数据集: {dataset_name}", flush=True)
+            print(f"输出目录: {eval_dir}", flush=True)
+            print("=" * 80, flush=True)
+
+        metric_tracker_Lpred, metric_tracker_Rpred, processing_times = evaluation_depth_custom_parallel(
+            rank,
+            world_size,
+            eval_dir,
+            test_dataset_config,
+            args,
+            pipeline,
+            test_data_dir,
+            alignment=alignment_type,
+            save_pred_vis=args.save_viz,
+        )
+
+        # 同步所有进程
+        dist.barrier()
+
+        # 收集两个CFG配置的结果
+        gathered_metrics_Lpred = [None] * world_size
+        gathered_metrics_Rpred = [None] * world_size
+        gathered_times = [None] * world_size
+
+        # 从metric_tracker获取结果字典
+        metrics_dict_Lpred = metric_tracker_Lpred.result() if hasattr(metric_tracker_Lpred, 'result') else {}
+        # metrics_dict_Rpred = metric_tracker_Rpred.result() if hasattr(metric_tracker_Rpred, 'result') else {}
+
+        dist.all_gather_object(gathered_metrics_Lpred, metrics_dict_Lpred)
+        # dist.all_gather_object(gathered_metrics_Rpred, metrics_dict_Rpred)
+        dist.all_gather_object(gathered_times, processing_times)
+
+        if rank == 0:
+            metrics_dict_Lpred, dataset_times = collect_and_merge_dual_cfg_results(rank, world_size, gathered_metrics_Lpred, gathered_times)
+
+            if metrics_dict_Lpred:
+                # 格式化并输出结果表格
+                formatted_output = format_dual_cfg_results_table(dataset_name, model_identifier, metrics_dict_Lpred, dataset_times)
+
+                print(formatted_output)
+                print(f"结果保存至: {eval_dir}")
+
+                # 存储结果用于后续汇总
+                all_dataset_results[dataset_name] = {'metrics_Lpred': metrics_dict_Lpred, 'formatted_output': formatted_output, 'eval_dir': eval_dir, 'processing_times': dataset_times}
+
+    if rank == 0:
+        print(f"\n所有数据集评估完成! 结果保存在: {os.path.join(args.output_dir, model_identifier)}")
+
+        # 保存所有数据集的双CFG评估结果摘要
+        if all_dataset_results:
+            save_dual_cfg_results_summary(os.path.join(args.output_dir, model_identifier), all_dataset_results, model_identifier)
+
+    cleanup()
+
+def main_worker_normal(rank, world_size, args, eval_datasets):
+    """normal预测的多进程主函数"""
+    from infer.inference import ImageGenerator
+    from infer.inner_evaluation import evaluation_normal_custom_parallel
+
+    setup(rank, world_size)
+    torch.cuda.set_device(rank)
+    device = torch.device(f"cuda:{rank}")
+
+    pipeline = ImageGenerator(
+        ae_path=os.path.join(args.model_path, 'vae.safetensors'),
+        dit_path=os.path.join(args.model_path, "step1x-edit-i1258-FP8.safetensors" if args.quantized else "step1x-edit-i1258.safetensors"),
+        qwen2vl_model_path=args.qwen2vl_model_path,
+        max_length=640,
+        quantized=args.quantized,
+        offload=args.offload,
+        lora=args.lora,
+        device=str(device),
+        args=args,
+    )
+
+    if rank == 0:
+        print(f"Successfully loading pipeline from {args.model_path}.")
+
+    test_data_dir = resolve_eval_data_root(args, "dsine_eval")
+    dataset_split_path = os.path.join(REPO_ROOT, "infer", "dataset_normal")
+    
+    # 使用新的模型标识符提取函数
+    model_identifier = extract_model_identifier(args.lora)
+    
+    # 修改输出目录结构：在任务名称外添加模型名称层级
+    eval_dir = os.path.join(args.output_dir, model_identifier, args.task_name)
+
+    if rank == 0:
+        print(f"模型标识符: {model_identifier}")
+        print(f"输出目录: {eval_dir}")
+
+    if rank == 0:
+        print(f"\n开始并行Normal评估，使用{world_size}个GPU")
+        print("=" * 80)
+
+    # 调用并行评估函数
+    all_normal_errors, all_processing_times, all_dataset_metrics = evaluation_normal_custom_parallel(
+        rank, world_size, eval_dir, test_data_dir, dataset_split_path, pipeline, args, eval_datasets, save_pred_vis=args.save_viz
+    )
+
+    # 同步所有进程
+    dist.barrier()
+
+    # 收集所有GPU的结果
+    gathered_normal_errors = [None] * world_size
+    gathered_processing_times = [None] * world_size
+    gathered_dataset_metrics = [None] * world_size
+
+    dist.all_gather_object(gathered_normal_errors, all_normal_errors)
+    dist.all_gather_object(gathered_processing_times, all_processing_times)
+    dist.all_gather_object(gathered_dataset_metrics, all_dataset_metrics)
+
+    if rank == 0:
+        # 合并所有GPU的结果
+        final_results = {}
+        
+        for dataset_name, _ in eval_datasets:
+            print(f"\n合并数据集 {dataset_name} 的结果...")
+            
+            # 合并normal errors
+            all_errors_for_dataset = []
+            all_times_for_dataset = []
+            
+            for gpu_errors, gpu_times in zip(gathered_normal_errors, gathered_processing_times):
+                if gpu_errors[dataset_name] is not None:
+                    all_errors_for_dataset.append(gpu_errors[dataset_name])
+                if gpu_times[dataset_name]:
+                    all_times_for_dataset.extend(gpu_times[dataset_name])
+            
+            # 计算最终指标
+            if all_errors_for_dataset:
+                combined_errors = torch.cat(all_errors_for_dataset, dim=0)
+                from infer.util import normal_utils
+                final_metrics = normal_utils.compute_normal_metrics(combined_errors)
+                
+                print(f"数据集 {dataset_name} 最终结果:")
+                print("mean median rmse 5 7.5 11.25 22.5 30")
+                print("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f" % (
+                    final_metrics['mean'], final_metrics['median'], final_metrics['rmse'], 
+                    final_metrics['a1'], final_metrics['a2'], final_metrics['a3'], 
+                    final_metrics['a4'], final_metrics['a5']
+                ))
+                
+                final_results[dataset_name] = {
+                    'metrics': final_metrics,
+                    'processing_times': all_times_for_dataset,
+                    'sample_count': len(combined_errors)
+                }
+                
+                # 保存结果到文件
+                dataset_output_dir = os.path.join(eval_dir, dataset_name)
+                os.makedirs(dataset_output_dir, exist_ok=True)
+                
+                from tabulate import tabulate
+                eval_text = f"Evaluation metrics for {dataset_name}:\n"
+                eval_text += f"Total samples: {len(combined_errors)}\n"
+                eval_text += f"Average processing time: {np.mean(all_times_for_dataset):.2f}s\n"
+                eval_text += tabulate([list(final_metrics.keys()), list(final_metrics.values())])
+                
+                save_path = os.path.join(dataset_output_dir, "eval_metrics.txt")
+                with open(save_path, "w+") as f:
+                    f.write(eval_text)
+                
+                print(f"结果已保存至: {save_path}")
+            else:
+                print(f"数据集 {dataset_name}: 未找到有效数据")
+                final_results[dataset_name] = None
+
+        # 保存总体结果摘要
+        summary_file = os.path.join(eval_dir, f"{model_identifier}_normal_summary.txt")
+        with open(summary_file, 'w', encoding='utf-8') as f:
+            f.write("=" * 120 + "\n")
+            f.write("Normal预测多GPU并行评估结果汇总\n")
+            f.write("=" * 120 + "\n\n")
+            
+            for dataset_name, result in final_results.items():
+                if result is not None:
+                    f.write(f"数据集: {dataset_name}\n")
+                    f.write(f"样本数: {result['sample_count']}\n")
+                    f.write(f"平均处理时间: {np.mean(result['processing_times']):.2f}s\n")
+                    metrics = result['metrics']
+                    f.write("mean median rmse 5 7.5 11.25 22.5 30\n")
+                    f.write("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n" % (
+                        metrics['mean'], metrics['median'], metrics['rmse'], 
+                        metrics['a1'], metrics['a2'], metrics['a3'], 
+                        metrics['a4'], metrics['a5']
+                    ))
+                    f.write("-" * 60 + "\n\n")
+                else:
+                    f.write(f"数据集: {dataset_name} - 无有效数据\n\n")
+        
+        print(f"\nNormal评估总结已保存至: {summary_file}")
+
+    cleanup()
+    
+def main():
+    args = parse_args()
+    if args.seed is not None:
+        seed_all(args.seed)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # 检查GPU数量
+    world_size = min(args.num_gpus, torch.cuda.device_count())
+    if world_size <= 0:
+        print("错误：未检测到可用的GPU。")
+        return
+
+    # 设置多进程相关环境变量
+    os.environ['OMP_NUM_THREADS'] = '1'
+    os.environ['MKL_NUM_THREADS'] = '1'
+
+    print(f"即将使用 {world_size} 个GPU进行并行推理...")
+
+    if args.task_name == 'depth':
+        test_depth_dataset_configs = parse_depth_eval_datasets(args.depth_eval_datasets)
+        print(f"Depth评估数据集: {list(test_depth_dataset_configs.keys())}")
+
+        if world_size == 1:
+            # 单GPU情况，直接运行
+            main_worker(0, 1, args, test_depth_dataset_configs)
+        else:
+            # 多GPU情况，使用multiprocessing
+            try:
+                mp.spawn(main_worker, args=(world_size, args, test_depth_dataset_configs), nprocs=world_size, join=True)
+            except Exception as e:
+                print(f"多进程执行出错: {e}")
+                print("尝试降级到单GPU模式...")
+                # 降级到单GPU模式
+                main_worker(0, 1, args, test_depth_dataset_configs)
+
+    elif args.task_name == 'normal':
+        eval_datasets = parse_normal_eval_datasets(args.normal_eval_datasets)
+        print(f"Normal评估数据集: {eval_datasets}")
+
+        if world_size == 1:
+            main_worker_normal(0, 1, args, eval_datasets)
+        else:
+            try:
+                mp.spawn(main_worker_normal, args=(world_size, args, eval_datasets), nprocs=world_size, join=True)
+            except Exception as e:
+                print(f"多进程执行出错: {e}")
+                print("尝试降级到单GPU模式...")
+                # 降级到单GPU模式
+                main_worker_normal(0, 1, args, eval_datasets)
+    else:
+        raise ValueError(f"不支持的 task_name: {args.task_name}，仅支持 depth 或 normal")
+
+if __name__ == '__main__':
+    main()

FE2E/infer/alignment.py

@@ -0,0 +1,70 @@
+import numpy as np
+import torch
+
+
+def align_depth_least_square(
+    gt_arr: np.ndarray,
+    pred_arr: np.ndarray,
+    valid_mask_arr: np.ndarray,
+    return_scale_shift=True,
+    max_resolution=None,
+):
+    ori_shape = pred_arr.shape  # input shape
+
+    gt = gt_arr.squeeze()  # [H, W]
+    pred = pred_arr.squeeze()
+    valid_mask = valid_mask_arr.squeeze()
+
+    # Downsample
+    if max_resolution is not None:
+        scale_factor = np.min(max_resolution / np.array(ori_shape[-2:]))
+        if scale_factor < 1:
+            downscaler = torch.nn.Upsample(scale_factor=scale_factor, mode="nearest")
+            gt = downscaler(torch.as_tensor(gt).unsqueeze(0)).numpy()
+            pred = downscaler(torch.as_tensor(pred).unsqueeze(0)).numpy()
+            valid_mask = (
+                downscaler(torch.as_tensor(valid_mask).unsqueeze(0).float())
+                .bool()
+                .numpy()
+            )
+
+    assert (
+        gt.shape == pred.shape == valid_mask.shape
+    ), f"{gt.shape}, {pred.shape}, {valid_mask.shape}"
+
+    gt_masked = gt[valid_mask].reshape((-1, 1))
+    pred_masked = pred[valid_mask].reshape((-1, 1))
+
+    # numpy solver
+    _ones = np.ones_like(pred_masked)
+    A = np.concatenate([pred_masked, _ones], axis=-1)
+    X = np.linalg.lstsq(A, gt_masked, rcond=None)[0]
+    scale, shift = X
+
+    aligned_pred = pred_arr * scale + shift
+
+    # restore dimensions
+    aligned_pred = aligned_pred.reshape(ori_shape)
+
+    if return_scale_shift:
+        return aligned_pred, scale, shift
+    else:
+        return aligned_pred
+
+
+# ******************** disparity space ********************
+def depth2disparity(depth, return_mask=False):
+    if isinstance(depth, torch.Tensor):
+        disparity = torch.zeros_like(depth)
+    elif isinstance(depth, np.ndarray):
+        disparity = np.zeros_like(depth)
+    non_negtive_mask = depth > 0
+    disparity[non_negtive_mask] = 1.0 / depth[non_negtive_mask]
+    if return_mask:
+        return disparity, non_negtive_mask
+    else:
+        return disparity
+
+
+def disparity2depth(disparity, **kwargs):
+    return depth2disparity(disparity, **kwargs)

FE2E/infer/configs/data_diode_all.yaml

@@ -0,0 +1,5 @@
+name: diode
+disp_name: diode_val_all
+dir: diode/diode_val.tar
+filenames: infer/data_split/diode/diode_val_all_filename_list.txt 
+processing_res: 768

FE2E/infer/configs/data_eth3d.yaml

@@ -0,0 +1,7 @@
+name: eth3d
+disp_name: eth3d_full
+# dir: eth3d
+dir: eth3d/eth3d.tar
+filenames: infer/data_split/eth3d/eth3d_filename_list.txt
+processing_res: 768
+alignment_max_res: 1024

FE2E/infer/configs/data_kitti_eigen_test.yaml

@@ -0,0 +1,6 @@
+name: kitti
+disp_name: kitti_eigen_test_full
+dir: kitti/kitti_eigen_split_test.tar
+filenames: infer/data_split/kitti/eigen_test_files_with_gt.txt
+kitti_bm_crop: true
+valid_mask_crop: eigen

FE2E/infer/configs/data_nyu_test.yaml

@@ -0,0 +1,5 @@
+name: nyu_v2
+disp_name: nyu_test_full
+dir: nyudepth/nyu_labeled_extracted.tar
+filenames: infer/data_split/nyu/filename_list_test.txt
+eigen_valid_mask: true

FE2E/infer/configs/data_scannet_val.yaml

@@ -0,0 +1,5 @@
+name: scannet
+disp_name: scannet_val_800_1
+# dir: scannet
+dir: scannet/scannet_val_sampled_800_1.tar
+filenames: infer/data_split/scannet/scannet_val_sampled_list_800_1.txt

FE2E/infer/dataset/__init__.py

@@ -0,0 +1,49 @@
+import os
+
+from .base_depth_dataset import BaseDepthDataset, get_pred_name, DatasetMode  # noqa: F401
+from .diode_dataset import DIODEDataset
+from .eth3d_dataset import ETH3DDataset
+from .kitti_dataset import KITTIDataset
+from .nyu_dataset import NYUDataset
+from .scannet_dataset import ScanNetDataset
+
+
+dataset_name_class_dict = {
+    "nyu_v2": NYUDataset,
+    "kitti": KITTIDataset,
+    "eth3d": ETH3DDataset,
+    "diode": DIODEDataset,
+    "scannet": ScanNetDataset,
+}
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def _resolve_split_file(path: str) -> str:
+    if os.path.isabs(path) or os.path.exists(path):
+        return path
+
+    candidate = os.path.join(REPO_ROOT, path)
+    if os.path.exists(candidate):
+        return candidate
+
+    return path
+
+
+def get_dataset(cfg_data_split, base_data_dir: str, mode: DatasetMode, prompt_type="query", **kwargs) -> BaseDepthDataset:
+    if cfg_data_split.name in dataset_name_class_dict.keys():
+        dataset_class = dataset_name_class_dict[cfg_data_split.name]
+        filename_ls_path = cfg_data_split.filenames if not prompt_type == "full" else (cfg_data_split.filenames).replace(".txt", "_wc.txt")
+        filename_ls_path = _resolve_split_file(filename_ls_path)
+        dataset = dataset_class(
+            mode=mode,
+            filename_ls_path=filename_ls_path,
+            dataset_dir=os.path.join(base_data_dir, cfg_data_split.dir),
+            **cfg_data_split,
+            prompt_type=prompt_type,
+            **kwargs,
+        )
+    else:
+        raise NotImplementedError
+
+    return dataset

FE2E/infer/dataset/base_depth_dataset.py

@@ -0,0 +1,257 @@
+import io
+import os
+import random
+import tarfile
+from enum import Enum
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision.transforms import InterpolationMode, Resize
+import torchvision.transforms.functional as F
+
+class DatasetMode(Enum):
+    RGB_ONLY = "rgb_only"
+    EVAL = "evaluate"
+    TRAIN = "train"
+
+
+def read_image_from_tar(tar_obj, img_rel_path):
+    image = tar_obj.extractfile("./" + img_rel_path)
+    image = image.read()
+    image = Image.open(io.BytesIO(image))
+
+
+class BaseDepthDataset(Dataset):
+
+    def __init__(
+        self, mode: DatasetMode, filename_ls_path: str,        dataset_dir: str,        disp_name: str,        min_depth,        max_depth,        has_filled_depth,        name_mode,        depth_transform=None,        augmentation_args: dict = None,        resize_to_hw=None,
+        move_invalid_to_far_plane: bool = True, rgb_transform=None, prompt_type="query", **kwargs,
+    ) -> None:
+        super().__init__()
+        self.mode = mode
+
+        self.filename_ls_path = filename_ls_path
+        self.dataset_dir = dataset_dir
+        self.disp_name = disp_name
+        self.has_filled_depth = has_filled_depth
+        self.name_mode: DepthFileNameMode = name_mode
+        self.min_depth = min_depth
+        self.max_depth = max_depth
+
+        self.depth_transform = depth_transform
+        self.augm_args = augmentation_args
+        self.resize_to_hw = resize_to_hw
+        self.prompt_type = prompt_type
+        # 设置默认的rgb_transform函数
+        if rgb_transform is None:
+            self.rgb_transform = self._default_rgb_transform
+        else:
+            self.rgb_transform = rgb_transform
+        self.move_invalid_to_far_plane = move_invalid_to_far_plane
+
+        # Load filenames
+        with open(self.filename_ls_path, "r") as f:
+            self.filenames = [s.split() for s in f.readlines()]  # [['rgb.png', 'depth.tif'], [], ...]
+
+        # Tar dataset
+        self.tar_obj = None
+        self.tar_obj_pid = None
+        self.is_tar = (True if os.path.isfile(dataset_dir) and tarfile.is_tarfile(dataset_dir) else False)
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, index):
+        rasters, other = self._get_data_item(index)
+        if DatasetMode.TRAIN == self.mode:
+            rasters = self._training_preprocess(rasters)
+        # merge
+        outputs = rasters
+        outputs.update(other)
+        return outputs
+
+    def _get_data_item(self, index):
+        rgb_rel_path, depth_rel_path, filled_rel_path, prompt = self._get_data_path(index=index)
+
+        rasters = {}
+
+        # RGB data
+        rasters.update(self._load_rgb_data(rgb_rel_path=rgb_rel_path))
+
+        # Depth data
+        if DatasetMode.RGB_ONLY != self.mode:
+            # load data
+            depth_data = self._load_depth_data(depth_rel_path=depth_rel_path, filled_rel_path=filled_rel_path)
+            rasters.update(depth_data)
+            # valid mask
+            rasters["valid_mask_raw"] = self._get_valid_mask(rasters["depth_raw_linear"]).clone()
+            rasters["valid_mask_filled"] = self._get_valid_mask(rasters["depth_filled_linear"]).clone()
+
+        
+        other = {"index": index, "rgb_relative_path": rgb_rel_path, "prompt": prompt}
+
+        return rasters, other
+
+    def _load_rgb_data(self, rgb_rel_path):
+        # Read RGB data
+        _, rgb = self._read_image(rgb_rel_path)
+        rgb = self.input_process_image(rgb)
+        outputs = {"rgb": rgb}
+        return outputs
+
+    def input_process_image(self, image):
+        if isinstance(image, np.ndarray):
+            image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
+            
+            return image
+        elif isinstance(image, Image.Image):
+            image = F.to_tensor(image.convert("RGB"))
+            
+            return image
+        elif isinstance(image, torch.Tensor):
+            return image
+        elif isinstance(image, str):
+            image = F.to_tensor(Image.open(image).convert("RGB"))
+            
+            return image
+        return image
+
+    def _load_depth_data(self, depth_rel_path, filled_rel_path):
+        # Read depth data
+        outputs = {}
+        depth_raw = self._read_depth_file(depth_rel_path).squeeze()
+        depth_raw_linear = torch.from_numpy(depth_raw).float().unsqueeze(0)  # [1, H, W]
+        outputs["depth_raw_linear"] = depth_raw_linear.clone()
+
+        if self.has_filled_depth:
+            depth_filled = self._read_depth_file(filled_rel_path).squeeze()
+            depth_filled_linear = torch.from_numpy(depth_filled).float().unsqueeze(0)
+            outputs["depth_filled_linear"] = depth_filled_linear
+        else:
+            outputs["depth_filled_linear"] = depth_raw_linear.clone()
+
+        return outputs
+
+    def _get_data_path(self, index):
+        filename_line = self.filenames[index]
+        rgb_rel_path = filename_line[0]
+        depth_rel_path, filled_rel_path = None, None
+        if DatasetMode.RGB_ONLY != self.mode:
+            depth_rel_path = filename_line[1]
+            if self.has_filled_depth:
+                filled_rel_path = filename_line[2]
+                
+        if self.prompt_type == "full":
+            if filename_line[2][0] in ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"]:
+                prompt = ' '.join(filename_line[2:])
+            else:
+                prompt = ' '.join(filename_line[3:])
+        else:
+            prompt = 1
+        return rgb_rel_path, depth_rel_path, filled_rel_path, prompt
+
+    def _read_image(self, img_rel_path):
+        if self.is_tar:
+            tar_obj = self._ensure_tar_obj()
+            image = tar_obj.extractfile("./" + img_rel_path)
+            image = image.read()
+            image = Image.open(io.BytesIO(image))
+        else:
+            img_path = os.path.join(self.dataset_dir, img_rel_path)
+            image = Image.open(img_path)
+        image_arr = np.asarray(image)
+        return image_arr, image
+
+    def _read_depth_file(self, rel_path):
+        depth_in, _ = self._read_image(rel_path)
+        #  Replace code below to decode depth according to dataset definition
+        depth_decoded = depth_in
+
+        return depth_decoded
+
+    def _get_valid_mask(self, depth: torch.Tensor):
+        valid_mask = torch.logical_and((depth > self.min_depth), (depth < self.max_depth)).bool()
+        return valid_mask
+
+    def _training_preprocess(self, rasters):
+        # Augmentation
+        if self.augm_args is not None:
+            rasters = self._augment_data(rasters)
+
+        # Normalization
+        rasters["depth_raw_norm"] = self.depth_transform(rasters["depth_raw_linear"], rasters["valid_mask_raw"]).clone()
+        rasters["depth_filled_norm"] = self.depth_transform(rasters["depth_filled_linear"], rasters["valid_mask_filled"]).clone()
+
+        # Set invalid pixel to far plane
+        if self.move_invalid_to_far_plane:
+            if self.depth_transform.far_plane_at_max:
+                rasters["depth_filled_norm"][~rasters["valid_mask_filled"]] = self.depth_transform.norm_max
+            else:
+                rasters["depth_filled_norm"][~rasters["valid_mask_filled"]] = self.depth_transform.norm_min
+
+        # Resize
+        if self.resize_to_hw is not None:
+            resize_transform = Resize(size=self.resize_to_hw, interpolation=InterpolationMode.NEAREST_EXACT)
+            rasters = {k: resize_transform(v) for k, v in rasters.items()}
+
+        return rasters
+
+    def _augment_data(self, rasters_dict):
+        # lr flipping
+        lr_flip_p = self.augm_args.lr_flip_p
+        if random.random() < lr_flip_p:
+            rasters_dict = {k: v.flip(-1) for k, v in rasters_dict.items()}
+
+        return rasters_dict
+
+    def __del__(self):
+        if self.tar_obj is not None:
+            self.tar_obj.close()
+            self.tar_obj = None
+            self.tar_obj_pid = None
+
+    def _default_rgb_transform(self, x):
+        """默认的RGB变换函数: [0, 255] -> [-1, 1]"""
+        return x / 255.0 * 2 - 1
+
+    def _ensure_tar_obj(self):
+        """Ensure each process owns its own tar handle to avoid cross-process FD issues."""
+        if not self.is_tar:
+            return None
+        current_pid = os.getpid()
+        if self.tar_obj is None or self.tar_obj_pid != current_pid:
+            if self.tar_obj is not None:
+                try:
+                    self.tar_obj.close()
+                except Exception:
+                    pass
+            self.tar_obj = tarfile.open(self.dataset_dir)
+            self.tar_obj_pid = current_pid
+        return self.tar_obj
+
+
+# Prediction file naming modes
+class DepthFileNameMode(Enum):
+    id = 1  # id.png
+    rgb_id = 2  # rgb_id.png
+    i_d_rgb = 3  # i_d_1_rgb.png
+    rgb_i_d = 4
+
+
+def get_pred_name(rgb_basename, name_mode, suffix=".png"):
+    if DepthFileNameMode.rgb_id == name_mode:
+        pred_basename = "pred_" + rgb_basename.split("_")[1]
+    elif DepthFileNameMode.i_d_rgb == name_mode:
+        pred_basename = rgb_basename.replace("_rgb.", "_pred.")
+    elif DepthFileNameMode.id == name_mode:
+        pred_basename = "pred_" + rgb_basename
+    elif DepthFileNameMode.rgb_i_d == name_mode:
+        pred_basename = "pred_" + "_".join(rgb_basename.split("_")[1:])
+    else:
+        raise NotImplementedError
+    # change suffix
+    pred_basename = os.path.splitext(pred_basename)[0] + suffix
+
+    return pred_basename

FE2E/infer/dataset/diode_dataset.py

@@ -0,0 +1,71 @@
+# Author: Bingxin Ke
+# Last modified: 2024-02-26
+
+import os
+import tarfile
+from io import BytesIO
+
+import numpy as np
+import torch
+
+from .base_depth_dataset import BaseDepthDataset, DepthFileNameMode, DatasetMode
+
+
+class DIODEDataset(BaseDepthDataset):
+    def __init__(
+        self,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            # DIODE data parameter
+            min_depth=0.6,
+            max_depth=350,
+            has_filled_depth=False,
+            name_mode=DepthFileNameMode.id,
+            **kwargs,
+        )
+
+    def _read_npy_file(self, rel_path):
+        if self.is_tar:
+            if self.tar_obj is None:
+                self.tar_obj = tarfile.open(self.dataset_dir)
+            fileobj = self.tar_obj.extractfile("./" + rel_path)
+            npy_path_or_content = BytesIO(fileobj.read())
+        else:
+            npy_path_or_content = os.path.join(self.dataset_dir, rel_path)
+        data = np.load(npy_path_or_content).squeeze()[np.newaxis, :, :]
+        return data
+
+    def _read_depth_file(self, rel_path):
+        depth = self._read_npy_file(rel_path)
+        return depth
+
+    def _get_data_path(self, index):
+        return self.filenames[index][0], self.filenames[index][1], self.filenames[index][2], self.filenames[index][3:] if len(self.filenames[index]) > 3 else 1
+
+    def _get_data_item(self, index):
+
+        rgb_rel_path, depth_rel_path, mask_rel_path, prompt = self._get_data_path(index=index)
+
+        rasters = {}
+
+        # RGB data
+        rasters.update(self._load_rgb_data(rgb_rel_path=rgb_rel_path))
+
+        # Depth data
+        if DatasetMode.RGB_ONLY != self.mode:
+            # load data
+            depth_data = self._load_depth_data(
+                depth_rel_path=depth_rel_path, filled_rel_path=None
+            )
+            rasters.update(depth_data)
+
+            # valid mask
+            mask = self._read_npy_file(mask_rel_path).astype(bool)
+            mask = torch.from_numpy(mask).bool()
+            rasters["valid_mask_raw"] = mask.clone()
+            rasters["valid_mask_filled"] = mask.clone()
+
+        other = {"index": index, "rgb_relative_path": rgb_rel_path, "prompt": prompt}
+
+        return rasters, other

FE2E/infer/dataset/drivingstereo_dataset.py

@@ -0,0 +1,32 @@
+import os
+import json
+import numpy as np
+import PIL.Image as pil
+
+from .mono_dataset import MonoDataset
+
+
+class DrivingStereoDataset(MonoDataset):
+    RAW_HEIGHT = 800
+    RAW_WIDTH = 1762
+
+    def __init__(self, *args, **kwargs):
+        super(DrivingStereoDataset, self).__init__(*args, **kwargs)
+        self.forename = {"rainy": "2018-08-17-09-45-58_2018-08-17-10-", "foggy": "2018-10-25-07-37-26_2018-10-25-", "sunny": "2018-10-19-09-30-39_2018-10-19-", "cloudy": "2018-10-31-06-55-01_2018-10-31-"}
+
+    def get_color(self, weather, name, do_flip):
+        path, name = self.get_image_path(weather, name)
+        color = self.loader(path)
+
+        return color, name
+
+    def get_image_path(self, weather, frame_name):
+        folder = "left-image-full-size"
+        image_path = os.path.join(self.opts.data_path, weather, folder, frame_name)
+        image_name = os.path.join(weather, folder, frame_name)
+        if self.opts.debug >= 3:
+            print(image_name)
+        return image_path, image_name
+
+    def index_to_name(self, weather, index):
+        return self.forename[weather] + self.filenames[index] + ".png"

FE2E/infer/dataset/eth3d_dataset.py

@@ -0,0 +1,45 @@
+# Author: Bingxin Ke
+# Last modified: 2024-02-08
+
+import torch
+import tarfile
+import os
+import numpy as np
+
+from .base_depth_dataset import BaseDepthDataset, DepthFileNameMode
+
+
+class ETH3DDataset(BaseDepthDataset):
+    HEIGHT, WIDTH = 4032, 6048
+
+    def __init__(
+        self,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            # ETH3D data parameter
+            min_depth=1e-5,
+            max_depth=torch.inf,
+            has_filled_depth=False,
+            name_mode=DepthFileNameMode.id,
+            **kwargs,
+        )
+
+    def _read_depth_file(self, rel_path):
+        # Read special binary data: https://www.eth3d.net/documentation#format-of-multi-view-data-image-formats
+        if self.is_tar:
+            tar_obj = self._ensure_tar_obj()
+            binary_data = tar_obj.extractfile("./" + rel_path)
+            binary_data = binary_data.read()
+
+        else:
+            depth_path = os.path.join(self.dataset_dir, rel_path)
+            with open(depth_path, "rb") as file:
+                binary_data = file.read()
+        # Convert the binary data to a numpy array of 32-bit floats
+        depth_decoded = np.frombuffer(binary_data, dtype=np.float32).copy()
+
+        depth_decoded[depth_decoded == torch.inf] = 0.0
+
+        depth_decoded = depth_decoded.reshape((self.HEIGHT, self.WIDTH))
+        return depth_decoded

FE2E/infer/dataset/kitti_dataset.py

@@ -0,0 +1,105 @@
+# Author: Bingxin Ke
+# Last modified: 2024-02-08
+
+import torch
+
+from .base_depth_dataset import BaseDepthDataset, DepthFileNameMode
+
+
+class KITTIDataset(BaseDepthDataset):
+    def __init__(
+            self,
+            kitti_bm_crop,  # Crop to KITTI benchmark size
+            valid_mask_crop,  # Evaluation mask. [None, garg or eigen]
+            **kwargs,
+    ) -> None:
+        super().__init__(
+            # KITTI data parameter
+            min_depth=1e-5,
+            max_depth=80,
+            has_filled_depth=False,
+            name_mode=DepthFileNameMode.id,
+            **kwargs,
+        )
+        self.kitti_bm_crop = kitti_bm_crop
+        self.valid_mask_crop = valid_mask_crop
+        assert self.valid_mask_crop in [
+            None,
+            "garg",  # set evaluation mask according to Garg  ECCV16
+            "eigen",  # set evaluation mask according to Eigen NIPS14
+        ], f"Unknown crop type: {self.valid_mask_crop}"
+
+        # Filter out empty depth
+        self.filenames = [f for f in self.filenames if "None" != f[1]]
+
+    def _read_depth_file(self, rel_path):
+        depth_in,_ = self._read_image(rel_path)
+        # Decode KITTI depth
+        depth_decoded = depth_in / 256.0
+        return depth_decoded
+
+    def _load_rgb_data(self, rgb_rel_path):
+        rgb_data = super()._load_rgb_data(rgb_rel_path)
+        if self.kitti_bm_crop:
+            rgb_data = {k: self.kitti_benchmark_crop(v) for k, v in rgb_data.items()}
+        return rgb_data
+
+    def _load_depth_data(self, depth_rel_path, filled_rel_path):
+        depth_data = super()._load_depth_data(depth_rel_path, filled_rel_path)
+        if self.kitti_bm_crop:
+            depth_data = {
+                k: self.kitti_benchmark_crop(v) for k, v in depth_data.items()
+            }
+        return depth_data
+
+    @staticmethod
+    def kitti_benchmark_crop(input_img):
+        """
+        Crop images to KITTI benchmark size
+        Args:
+            `input_img` (torch.Tensor): Input image to be cropped.
+
+        Returns:
+            torch.Tensor:Cropped image.
+        """
+        KB_CROP_HEIGHT = 352
+        KB_CROP_WIDTH = 1216
+
+        height, width = input_img.shape[-2:]
+        top_margin = int(height - KB_CROP_HEIGHT)
+        left_margin = int((width - KB_CROP_WIDTH) / 2)
+        if 2 == len(input_img.shape):
+            out = input_img[
+                  top_margin: top_margin + KB_CROP_HEIGHT,
+                  left_margin: left_margin + KB_CROP_WIDTH,
+                  ]
+        elif 3 == len(input_img.shape):
+            out = input_img[
+                  :,
+                  top_margin: top_margin + KB_CROP_HEIGHT,
+                  left_margin: left_margin + KB_CROP_WIDTH,
+                  ]
+        return out
+
+    def _get_valid_mask(self, depth: torch.Tensor):
+        # reference: https://github.com/cleinc/bts/blob/master/pytorch/bts_eval.py
+        valid_mask = super()._get_valid_mask(depth)  # [1, H, W]
+
+        if self.valid_mask_crop is not None:
+            eval_mask = torch.zeros_like(valid_mask.squeeze()).bool()
+            gt_height, gt_width = eval_mask.shape
+
+            if "garg" == self.valid_mask_crop:
+                eval_mask[
+                int(0.40810811 * gt_height): int(0.99189189 * gt_height),
+                int(0.03594771 * gt_width): int(0.96405229 * gt_width),
+                ] = 1
+            elif "eigen" == self.valid_mask_crop:
+                eval_mask[
+                int(0.3324324 * gt_height): int(0.91351351 * gt_height),
+                int(0.0359477 * gt_width): int(0.96405229 * gt_width),
+                ] = 1
+
+            eval_mask.reshape(valid_mask.shape)
+            valid_mask = torch.logical_and(valid_mask, eval_mask)
+        return valid_mask

FE2E/infer/dataset/nyu_dataset.py

@@ -0,0 +1,43 @@
+# Author: Bingxin Ke
+# Last modified: 2024-02-08
+
+
+import torch
+
+from .base_depth_dataset import BaseDepthDataset, DepthFileNameMode
+
+
+class NYUDataset(BaseDepthDataset):
+    def __init__(
+        self,
+        eigen_valid_mask: bool,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            # NYUv2 dataset parameter
+            min_depth=1e-3,
+            max_depth=10.0,
+            has_filled_depth=True,
+            name_mode=DepthFileNameMode.rgb_id,
+            **kwargs,
+        )
+
+        self.eigen_valid_mask = eigen_valid_mask
+
+    def _read_depth_file(self, rel_path):
+        depth_in,_ = self._read_image(rel_path)
+        # Decode NYU depth
+        depth_decoded = depth_in / 1000.0
+        return depth_decoded
+
+    def _get_valid_mask(self, depth: torch.Tensor):
+        valid_mask = super()._get_valid_mask(depth)
+
+        # Eigen crop for evaluation
+        if self.eigen_valid_mask:
+            eval_mask = torch.zeros_like(valid_mask.squeeze()).bool()
+            eval_mask[45:471, 41:601] = 1
+            eval_mask.reshape(valid_mask.shape)
+            valid_mask = torch.logical_and(valid_mask, eval_mask)
+
+        return valid_mask

FE2E/infer/dataset/scannet_dataset.py

@@ -0,0 +1,25 @@
+# Author: Bingxin Ke
+# Last modified: 2024-02-08
+
+from .base_depth_dataset import BaseDepthDataset, DepthFileNameMode
+
+
+class ScanNetDataset(BaseDepthDataset):
+    def __init__(
+        self,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            # ScanNet data parameter
+            min_depth=1e-3,
+            max_depth=10,
+            has_filled_depth=False,
+            name_mode=DepthFileNameMode.id,
+            **kwargs,
+        )
+
+    def _read_depth_file(self, rel_path):
+        depth_in,_ = self._read_image(rel_path)
+        # Decode ScanNet depth
+        depth_decoded = depth_in / 1000.0
+        return depth_decoded

FE2E/infer/dataset_normal/__init__.py

@@ -0,0 +1,27 @@
+""" data sample
+"""
+class Sample():
+    def __init__(self, img=None, 
+                 depth=None, depth_mask=None,
+                 normal=None, normal_mask=None, 
+                 intrins=None, flipped=False,
+                 dataset_name='dataset', scene_name='scene', img_name='img', 
+                 info={}):
+ 
+        self.img = img                  # input image
+
+        self.depth = depth              # depth - GT
+        self.depth_mask = depth_mask    # depth - valid_mask
+
+        self.normal = normal            # surface normals - GT
+        self.normal_mask = normal_mask  # surface normals - valid_mask
+        
+        self.intrins = intrins          # camera intrinsics
+        self.flipped = flipped          # True when the image is flipped during augmentation
+
+        self.dataset_name = dataset_name
+        self.scene_name = scene_name
+        self.img_name = img_name
+
+        # other info (this is a dict containing any additional information)
+        self.info = info

FE2E/infer/dataset_normal/aug_basic.py

@@ -0,0 +1,239 @@
+""" basic augmentations
+"""
+import random
+import numpy as np
+
+import torch
+from torchvision import transforms
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+
+import logging
+logger = logging.getLogger('root')
+
+
+def resize(sample, new_H, new_W):
+    _, orig_H, orig_W = sample.img.shape
+    sample.img = F.interpolate(sample.img.unsqueeze(0), size=(new_H, new_W), mode='bilinear', align_corners=False, antialias=True).squeeze(0)
+    if sample.depth is not None:
+        sample.depth = F.interpolate(sample.depth.unsqueeze(0), size=(new_H, new_W), mode='nearest').squeeze(0)
+    if sample.depth_mask is not None:
+        sample.depth_mask = F.interpolate(sample.depth_mask.unsqueeze(0).float(), size=(new_H, new_W), mode='nearest').squeeze(0) > 0.5
+    if sample.normal is not None:
+        sample.normal = F.interpolate(sample.normal.unsqueeze(0), size=(new_H, new_W), mode='nearest').squeeze(0)
+    if sample.normal_mask is not None:
+        sample.normal_mask = F.interpolate(sample.normal_mask.unsqueeze(0).float(), size=(new_H, new_W), mode='nearest').squeeze(0) > 0.5
+    if sample.intrins is not None:
+        # NOTE: top-left is (0,0)
+        sample.intrins[0, 0] = sample.intrins[0, 0] * (new_W / orig_W)     # fx
+        sample.intrins[1, 1] = sample.intrins[1, 1] * (new_H / orig_H)     # fy
+        sample.intrins[0, 2] = (sample.intrins[0, 2] + 0.5) * (new_W / orig_W) - 0.5     # cx
+        sample.intrins[1, 2] = (sample.intrins[1, 2] + 0.5) * (new_H / orig_H) - 0.5     # cy
+    return sample
+
+
+def pad(sample, lrtb):
+    l, r, t, b = lrtb
+    sample.img = F.pad(sample.img, (l, r, t, b), mode="constant", value=0)
+    if sample.depth is not None:
+        sample.depth = F.pad(sample.depth, (l, r, t, b), mode="constant", value=0)
+    if sample.depth_mask is not None:
+        sample.depth_mask = F.pad(sample.depth_mask, (l, r, t, b), mode="constant", value=False)
+    if sample.normal is not None:
+        sample.normal = F.pad(sample.normal, (l, r, t, b), mode="constant", value=0)
+    if sample.normal_mask is not None:
+        sample.normal_mask = F.pad(sample.normal_mask, (l, r, t, b), mode="constant", value=False)
+    if sample.intrins is not None:
+        sample.intrins[0, 2] = sample.intrins[0, 2] + l
+        sample.intrins[1, 2] = sample.intrins[1, 2] + t
+    return sample
+
+
+def crop(sample, y, H, x, W):
+    sample.img = sample.img[:, y:y+H, x:x+W]
+    if sample.depth is not None:
+        sample.depth = sample.depth[:, y:y+H, x:x+W]
+    if sample.depth_mask is not None:
+        sample.depth_mask = sample.depth_mask[:, y:y+H, x:x+W]
+    if sample.normal is not None:
+        sample.normal = sample.normal[:, y:y+H, x:x+W]
+    if sample.normal_mask is not None:
+        sample.normal_mask = sample.normal_mask[:, y:y+H, x:x+W]
+    if sample.intrins is not None:
+        sample.intrins[0, 2] = sample.intrins[0, 2] - x
+        sample.intrins[1, 2] = sample.intrins[1, 2] - y
+    return sample
+
+
+class ToTensor():
+    """ numpy arrays to torch tensors
+    """
+    def __call__(self, sample):
+        sample.img = torch.from_numpy(sample.img).permute(2, 0, 1)                      # (3, H, W)
+        if sample.depth is not None:
+            sample.depth = torch.from_numpy(sample.depth).permute(2, 0, 1)              # (1, H, W)
+        if sample.depth_mask is not None:
+            sample.depth_mask = torch.from_numpy(sample.depth_mask).permute(2, 0, 1)    # (1, H, W)
+        if sample.normal is not None:
+            sample.normal = torch.from_numpy(sample.normal).permute(2, 0, 1)            # (3, H, W)
+        if sample.normal_mask is not None:
+            sample.normal_mask = torch.from_numpy(sample.normal_mask).permute(2, 0, 1)  # (1, H, W)
+        if sample.intrins is not None:
+            sample.intrins = torch.from_numpy(sample.intrins)                           # (3, 3)
+        return sample
+
+
+class RandomIntrins():
+    """ randomize intrinsics
+        sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1]
+    """
+    def __call__(self, sample):
+        assert 'crop_H' in sample.info.keys() 
+        assert 'crop_W' in sample.info.keys()
+        crop_H = sample.info['crop_H']
+        crop_W = sample.info['crop_W']
+
+        # height-based resizing
+        _, orig_H, orig_W = sample.img.shape
+        new_H = random.randrange(min(orig_H, crop_H), max(orig_H, crop_H)+1)
+        new_W = round((new_H / orig_H) * orig_W)
+        sample = resize(sample, new_H=new_H, new_W=new_W)
+
+        # pad if necessary
+        orig_H, orig_W = sample.img.shape[1], sample.img.shape[2]
+        l, r, t, b = 0, 0, 0, 0
+        if crop_H > orig_H:
+            t = b = crop_H - orig_H
+        if crop_W > orig_W:
+            l = r = crop_W - orig_W
+        sample = pad(sample, (l, r, t, b))
+
+        # crop
+        assert sample.img.shape[1] >= crop_H
+        assert sample.img.shape[2] >= crop_W
+        x = random.randint(0, sample.img.shape[2] - crop_W)
+        y = random.randint(0, sample.img.shape[1] - crop_H)
+        sample = crop(sample, y=y, H=crop_H, x=x, W=crop_W)
+
+        return sample
+
+
+class Resize():
+    """ resize to (H, W)
+        sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1]
+    """
+    def __init__(self, H=480, W=640):
+        self.H = H
+        self.W = W
+
+    def __call__(self, sample):
+        return resize(sample, new_H=self.H, new_W=self.W)
+
+
+class RandomCrop():
+    """ random crop
+        sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1]
+    """
+    def __init__(self, H=416, W=544):
+        self.H = H
+        self.W = W
+
+    def __call__(self, sample):
+        assert sample.img.shape[1] >= self.H
+        assert sample.img.shape[2] >= self.W
+        x = random.randint(0, sample.img.shape[2] - self.W)
+        y = random.randint(0, sample.img.shape[1] - self.H)
+        return crop(sample, y=y, H=self.H, x=x, W=self.W)
+    
+
+class NyuCrop():
+    """ crop image border for NYUv2 images
+        W = 43:608 / H = 45:472
+        sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1]
+    """
+    def __call__(self, sample):
+        return crop(sample, y=45, H=472-45, x=43, W=608-43)
+
+
+class HorizontalFlip():
+    """ random horizontal flipping
+        sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1]
+    """
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, sample):
+        if random.random() < self.p:
+            sample.img = TF.hflip(sample.img)
+            if sample.depth is not None:
+                sample.depth = TF.hflip(sample.depth)
+            if sample.depth_mask is not None:
+                sample.depth_mask = TF.hflip(sample.depth_mask)
+            if sample.normal is not None:
+                sample.normal = TF.hflip(sample.normal)
+                sample.normal[0, :, :] = -sample.normal[0, :, :]
+            if sample.normal_mask is not None:
+                sample.normal_mask = TF.hflip(sample.normal_mask)
+            if sample.intrins is not None:
+                # NOTE: top-left is (0,0)
+                _, H, W = sample.img.shape
+                sample.intrins[0, 2] = sample.intrins[0, 2] + 0.5   # top-left is (0.5, 0.5)
+                sample.intrins[0, 2] = W - sample.intrins[0, 2]
+                sample.intrins[0, 2] = sample.intrins[0, 2] - 0.5   # top-left is (0, 0)
+            sample.flipped = True
+        return sample
+
+
+class ColorAugmentation():
+    """ color augmentation
+        sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1]
+    """
+    def __init__(self, gamma_range=(0.9, 1.1),
+                       brightness_range=(0.75, 1.25),
+                       color_range=(0.9, 1.1),
+                       p=0.5):
+        self.gamma_range = gamma_range
+        self.brightness_range = brightness_range
+        self.color_range = color_range
+        self.p = p
+
+    def __call__(self, sample):
+        if random.random() < self.p:
+            # gamma augmentation
+            gamma = random.uniform(*self.gamma_range)
+            sample.img = sample.img ** gamma
+
+            # brightness augmentation
+            brightness = random.uniform(*self.brightness_range)
+            sample.img = sample.img * brightness
+
+            # color augmentation
+            colors = np.random.uniform(*self.color_range, size=3).astype(np.float32)
+            colors = torch.from_numpy(colors).view(3, 1, 1)
+            sample.img = sample.img * colors
+
+            # clip
+            sample.img = torch.clip(sample.img, 0, 1)
+
+        return sample
+    
+
+class Normalize():
+    """ mean & std: for image normalization
+        sample.img is a torch tensor of shape (3, H, W), normalized to [0, 1]
+    """
+    def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+        self.normalize = transforms.Normalize(mean=mean, std=std)
+
+    def __call__(self, sample):
+        sample.img = self.normalize(torch.clip(sample.img, min=0.0, max=1.0))
+        return sample
+
+
+class ToDict():
+    def __call__(self, sample):
+        data_dict = {}
+        for k, v in vars(sample).items():
+            if v is not None:
+                data_dict[k] = v
+        return data_dict

FE2E/infer/dataset_normal/hypersim/__init__.py

@@ -0,0 +1,65 @@
+""" Get samples from Hypersim dataset
+    Based on vkitti implementation
+"""
+import os
+import cv2
+import numpy as np
+
+from infer.dataset_normal import Sample
+
+
+def get_sample(base_data_dir, sample_path, info):
+    # e.g. sample_path = "ai_001_001/rgb_cam_00_fr0000.png"
+    scene_name = sample_path.split('/')[0]
+    img_filename = sample_path.split('/')[1]
+    
+    # Extract frame number from filename like "rgb_cam_00_fr0000.png"
+    frame_num = img_filename.split('_fr')[1].split('.')[0]
+    
+    dataset_path = os.path.join(base_data_dir, 'dsine_eval', 'hypersim')
+    img_path = os.path.join(dataset_path, sample_path)
+    
+    # Build corresponding normal path
+    normal_filename = f'depth_plane_cam_00_fr{frame_num}_1024x0768_normal_decoded_normal.png'
+    normal_path = os.path.join(dataset_path, scene_name, normal_filename)
+    
+    assert os.path.exists(img_path), f"Image not found: {img_path}"
+    assert os.path.exists(normal_path), f"Normal not found: {normal_path}"
+
+    # read image (H, W, 3)
+    img = cv2.cvtColor(cv2.imread(img_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32) / 255.0
+    
+    # read normal (H, W, 3)
+    normal = cv2.cvtColor(cv2.imread(normal_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    normal_mask = np.sum(normal, axis=2, keepdims=True) > 0
+    normal = (normal.astype(np.float32) / 255.0) * 2.0 - 1.0
+
+    # Create default intrinsics since hypersim doesn't provide them
+    # Using typical values for 1024x768 resolution
+    H, W = img.shape[:2]
+    fx = fy = 0.8 * W  # Typical focal length assumption
+    cx = W / 2.0
+    cy = H / 2.0
+    intrins = np.array([
+        [fx, 0, cx],
+        [0, fy, cy],
+        [0, 0, 1]
+    ], dtype=np.float32)
+
+    # Extract img_name for compatibility
+    img_name = img_filename.replace('.png', '')
+
+    sample = Sample(
+        img=img,
+        normal=normal,
+        normal_mask=normal_mask,
+        intrins=intrins,
+
+        dataset_name='hypersim',
+        scene_name=scene_name,
+        img_name=img_name,
+        info=info
+    )
+
+    return sample 

FE2E/infer/dataset_normal/hypersim/split/hypersim.txt

@@ -0,0 +1,98 @@
+ai_001_001/rgb_cam_00_fr0000.png
+ai_001_001/rgb_cam_00_fr0001.png
+ai_001_001/rgb_cam_00_fr0002.png
+ai_001_001/rgb_cam_00_fr0003.png
+ai_001_001/rgb_cam_00_fr0004.png
+ai_001_001/rgb_cam_00_fr0005.png
+ai_001_001/rgb_cam_00_fr0006.png
+ai_001_001/rgb_cam_00_fr0007.png
+ai_001_001/rgb_cam_00_fr0008.png
+ai_001_001/rgb_cam_00_fr0009.png
+ai_001_001/rgb_cam_00_fr0010.png
+ai_001_001/rgb_cam_00_fr0011.png
+ai_001_001/rgb_cam_00_fr0012.png
+ai_001_001/rgb_cam_00_fr0013.png
+ai_001_001/rgb_cam_00_fr0014.png
+ai_001_001/rgb_cam_00_fr0015.png
+ai_001_001/rgb_cam_00_fr0016.png
+ai_001_001/rgb_cam_00_fr0017.png
+ai_001_001/rgb_cam_00_fr0018.png
+ai_001_001/rgb_cam_00_fr0019.png
+ai_001_001/rgb_cam_00_fr0020.png
+ai_001_001/rgb_cam_00_fr0021.png
+ai_001_001/rgb_cam_00_fr0022.png
+ai_001_001/rgb_cam_00_fr0023.png
+ai_001_001/rgb_cam_00_fr0024.png
+ai_001_001/rgb_cam_00_fr0025.png
+ai_001_001/rgb_cam_00_fr0026.png
+ai_001_001/rgb_cam_00_fr0027.png
+ai_001_001/rgb_cam_00_fr0028.png
+ai_001_001/rgb_cam_00_fr0029.png
+ai_001_001/rgb_cam_00_fr0030.png
+ai_001_001/rgb_cam_00_fr0031.png
+ai_001_001/rgb_cam_00_fr0032.png
+ai_001_001/rgb_cam_00_fr0033.png
+ai_001_001/rgb_cam_00_fr0034.png
+ai_001_001/rgb_cam_00_fr0035.png
+ai_001_001/rgb_cam_00_fr0037.png
+ai_001_001/rgb_cam_00_fr0038.png
+ai_001_001/rgb_cam_00_fr0039.png
+ai_001_001/rgb_cam_00_fr0040.png
+ai_001_001/rgb_cam_00_fr0041.png
+ai_001_001/rgb_cam_00_fr0042.png
+ai_001_001/rgb_cam_00_fr0043.png
+ai_001_001/rgb_cam_00_fr0044.png
+ai_001_001/rgb_cam_00_fr0045.png
+ai_001_001/rgb_cam_00_fr0046.png
+ai_001_001/rgb_cam_00_fr0047.png
+ai_001_001/rgb_cam_00_fr0048.png
+ai_001_001/rgb_cam_00_fr0049.png
+ai_001_001/rgb_cam_00_fr0050.png
+ai_001_001/rgb_cam_00_fr0051.png
+ai_001_001/rgb_cam_00_fr0052.png
+ai_001_001/rgb_cam_00_fr0053.png
+ai_001_001/rgb_cam_00_fr0054.png
+ai_001_001/rgb_cam_00_fr0055.png
+ai_001_001/rgb_cam_00_fr0056.png
+ai_001_001/rgb_cam_00_fr0057.png
+ai_001_001/rgb_cam_00_fr0058.png
+ai_001_001/rgb_cam_00_fr0059.png
+ai_001_001/rgb_cam_00_fr0060.png
+ai_001_001/rgb_cam_00_fr0062.png
+ai_001_001/rgb_cam_00_fr0063.png
+ai_001_001/rgb_cam_00_fr0064.png
+ai_001_001/rgb_cam_00_fr0065.png
+ai_001_001/rgb_cam_00_fr0066.png
+ai_001_001/rgb_cam_00_fr0067.png
+ai_001_001/rgb_cam_00_fr0068.png
+ai_001_001/rgb_cam_00_fr0069.png
+ai_001_001/rgb_cam_00_fr0070.png
+ai_001_001/rgb_cam_00_fr0071.png
+ai_001_001/rgb_cam_00_fr0072.png
+ai_001_001/rgb_cam_00_fr0073.png
+ai_001_001/rgb_cam_00_fr0074.png
+ai_001_001/rgb_cam_00_fr0075.png
+ai_001_001/rgb_cam_00_fr0076.png
+ai_001_001/rgb_cam_00_fr0077.png
+ai_001_001/rgb_cam_00_fr0078.png
+ai_001_001/rgb_cam_00_fr0079.png
+ai_001_001/rgb_cam_00_fr0080.png
+ai_001_001/rgb_cam_00_fr0081.png
+ai_001_001/rgb_cam_00_fr0082.png
+ai_001_001/rgb_cam_00_fr0083.png
+ai_001_001/rgb_cam_00_fr0084.png
+ai_001_001/rgb_cam_00_fr0085.png
+ai_001_001/rgb_cam_00_fr0086.png
+ai_001_001/rgb_cam_00_fr0087.png
+ai_001_001/rgb_cam_00_fr0088.png
+ai_001_001/rgb_cam_00_fr0089.png
+ai_001_001/rgb_cam_00_fr0090.png
+ai_001_001/rgb_cam_00_fr0091.png
+ai_001_001/rgb_cam_00_fr0092.png
+ai_001_001/rgb_cam_00_fr0093.png
+ai_001_001/rgb_cam_00_fr0094.png
+ai_001_001/rgb_cam_00_fr0095.png
+ai_001_001/rgb_cam_00_fr0096.png
+ai_001_001/rgb_cam_00_fr0097.png
+ai_001_001/rgb_cam_00_fr0098.png
+ai_001_001/rgb_cam_00_fr0099.png

FE2E/infer/dataset_normal/ibims/__init__.py

@@ -0,0 +1,47 @@
+""" Get samples from iBims-1 (https://paperswithcode.com/dataset/ibims-1)
+    NOTE: We computed the GT surface normals by doing discontinuity-aware plane fitting
+"""
+import os
+import cv2
+import numpy as np
+os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
+
+from infer.dataset_normal import Sample
+
+def get_sample(base_data_dir, sample_path, info):
+    # e.g. sample_path = "ibims/corridor_01_img.png"
+    scene_name = sample_path.split('/')[0]
+    img_name, img_ext = sample_path.split('/')[1].split('_img')
+
+    dataset_path = os.path.join(base_data_dir, 'dsine_eval', 'ibims')
+    img_path = '%s/%s' % (dataset_path, sample_path)
+    normal_path = img_path.replace('_img'+img_ext, '_normal.exr')
+    intrins_path = img_path.replace('_img'+img_ext, '_intrins.npy')
+    assert os.path.exists(img_path)
+    assert os.path.exists(normal_path)
+    assert os.path.exists(intrins_path)
+
+    # read image (H, W, 3)
+    img = cv2.cvtColor(cv2.imread(img_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32) / 255.0
+
+    # read normal (H, W, 3)
+    normal = cv2.cvtColor(cv2.imread(normal_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    normal_mask = np.linalg.norm(normal, axis=2, keepdims=True) > 0.5
+
+    # read intrins (3, 3)
+    intrins = np.load(intrins_path)
+
+    sample = Sample(
+        img=img,
+        normal=normal,
+        normal_mask=normal_mask,
+        intrins=intrins,
+
+        dataset_name='ibims',
+        scene_name=scene_name,
+        img_name=img_name,
+        info=info
+    )
+
+    return sample

FE2E/infer/dataset_normal/ibims/split/ibims.txt

@@ -0,0 +1,100 @@
+ibims/corridor_01_img.png
+ibims/corridor_02_img.png
+ibims/corridor_03_img.png
+ibims/corridor_04_img.png
+ibims/corridor_05_img.png
+ibims/corridor_06_img.png
+ibims/corridor_07_img.png
+ibims/corridor_08_img.png
+ibims/corridor_09_img.png
+ibims/corridor_10_img.png
+ibims/factory_01_img.png
+ibims/factory_02_img.png
+ibims/factory_03_img.png
+ibims/factory_04_img.png
+ibims/factory_05_img.png
+ibims/factory_06_img.png
+ibims/factory_07_img.png
+ibims/factory_08_img.png
+ibims/kitchen_01_img.png
+ibims/kitchen_02_img.png
+ibims/kitchen_03_img.png
+ibims/kitchen_04_img.png
+ibims/kitchen_05_img.png
+ibims/kitchen_06_img.png
+ibims/kitchen_07_img.png
+ibims/kitchen_08_img.png
+ibims/lab_01_img.png
+ibims/lab_02_img.png
+ibims/lab_03_img.png
+ibims/lab_04_img.png
+ibims/lab_05_img.png
+ibims/lab_06_img.png
+ibims/lab_07_img.png
+ibims/lab_08_img.png
+ibims/lab_09_img.png
+ibims/lab_10_img.png
+ibims/lab_11_img.png
+ibims/lectureroom_01_img.png
+ibims/lectureroom_02_img.png
+ibims/lectureroom_03_img.png
+ibims/lectureroom_04_img.png
+ibims/lectureroom_05_img.png
+ibims/lectureroom_06_img.png
+ibims/lectureroom_07_img.png
+ibims/lectureroom_08_img.png
+ibims/lectureroom_09_img.png
+ibims/lectureroom_10_img.png
+ibims/livingroom_01_img.png
+ibims/livingroom_02_img.png
+ibims/livingroom_03_img.png
+ibims/livingroom_04_img.png
+ibims/livingroom_05_img.png
+ibims/livingroom_06_img.png
+ibims/livingroom_07_img.png
+ibims/livingroom_08_img.png
+ibims/livingroom_09_img.png
+ibims/livingroom_10_img.png
+ibims/livingroom_11_img.png
+ibims/livingroom_12_img.png
+ibims/livingroom_13_img.png
+ibims/livingroom_14_img.png
+ibims/livingroom_15_img.png
+ibims/meetingroom_01_img.png
+ibims/meetingroom_02_img.png
+ibims/meetingroom_03_img.png
+ibims/meetingroom_04_img.png
+ibims/meetingroom_05_img.png
+ibims/meetingroom_06_img.png
+ibims/meetingroom_07_img.png
+ibims/meetingroom_08_img.png
+ibims/office_01_img.png
+ibims/office_02_img.png
+ibims/office_03_img.png
+ibims/office_04_img.png
+ibims/office_05_img.png
+ibims/office_06_img.png
+ibims/office_07_img.png
+ibims/office_08_img.png
+ibims/restaurant_01_img.png
+ibims/restaurant_02_img.png
+ibims/restaurant_03_img.png
+ibims/restaurant_04_img.png
+ibims/restaurant_05_img.png
+ibims/restaurant_06_img.png
+ibims/restaurant_07_img.png
+ibims/restaurant_08_img.png
+ibims/restaurant_09_img.png
+ibims/restaurant_10_img.png
+ibims/restaurant_11_img.png
+ibims/restaurant_12_img.png
+ibims/restroom_01_img.png
+ibims/restroom_02_img.png
+ibims/storageroom_01_img.png
+ibims/storageroom_02_img.png
+ibims/storageroom_03_img.png
+ibims/storageroom_04_img.png
+ibims/storageroom_05_img.png
+ibims/storageroom_06_img.png
+ibims/storageroom_07_img.png
+ibims/storageroom_08_img.png

FE2E/infer/dataset_normal/normal_dataloader.py

@@ -0,0 +1,83 @@
+import os
+import random
+
+import torch
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+from . import aug_basic
+
+import logging
+logger = logging.getLogger('root')
+
+
+def get_transform(dataset_name='hypersim', mode='test'):
+    assert mode in ['test']
+    logger.info('Defining %s transform for %s dataset' % (mode, dataset_name))
+    tf_list = [
+        aug_basic.ToTensor(),
+    ]
+    tf_list += [
+        # 选项1：使用标准归一化 (如果启用，需要相应调整可视化代码)
+        # aug_basic.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        # 选项2：不使用归一化 (当前方案，配合修改后的unnormalize函数)
+        aug_basic.ToDict(),
+    ]
+    logger.info('Defining %s transform for %s dataset ... DONE' % (mode, dataset_name))
+    return transforms.Compose(tf_list)
+
+
+class NormalDataset(Dataset):
+    def __init__(self, base_data_dir, dataset_split_path, dataset_name='nyuv2', split='test', mode='test',  epoch=0):
+        self.split = split
+        self.mode = mode
+        self.base_data_dir = base_data_dir 
+        assert mode in ['test']
+
+        # data split
+        split_path = os.path.join(dataset_split_path, dataset_name, 'split', split+'.txt') # dataset_split_path: eval/dataset_normal/
+        assert os.path.exists(split_path)
+        with open(split_path, 'r') as f:
+            self.filenames = [i.strip() for i in f.readlines()]
+        self.split_path = split_path
+
+        # get_sample function
+        if dataset_name == 'nyuv2':
+            from infer.dataset_normal.nyuv2 import get_sample
+        elif dataset_name == 'scannet':
+            from infer.dataset_normal.scannet import get_sample
+        elif dataset_name == 'ibims':
+            from infer.dataset_normal.ibims import get_sample
+        elif dataset_name == 'sintel':
+            from infer.dataset_normal.sintel import get_sample
+        elif dataset_name == 'oasis':
+            from infer.dataset_normal.oasis import get_sample
+        elif dataset_name == 'hypersim':
+            from infer.dataset_normal.hypersim import get_sample
+        else:
+            raise NotImplementedError(f"Unsupported normal dataset: {dataset_name}")
+        self.get_sample = get_sample
+
+        # data preprocessing/augmentation
+        self.transform = get_transform(dataset_name=dataset_name, mode=mode)
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, index):
+        info = {}
+
+        sample = self.transform(self.get_sample(
+            base_data_dir = self.base_data_dir,
+            sample_path=self.filenames[index], 
+            info=info)
+        )
+
+        return sample            
+    
+class TestLoader(object):
+    def __init__(self, base_data_dir, dataset_split_path, dataset_name_test, test_split):
+        self.test_samples = NormalDataset(base_data_dir, dataset_split_path, dataset_name=dataset_name_test, 
+                                          split=test_split, mode='test', epoch=0)
+        self.data = DataLoader(self.test_samples, 1, shuffle=False, num_workers=4, pin_memory=True)

FE2E/infer/dataset_normal/nyuv2/__init__.py

@@ -0,0 +1,66 @@
+""" Get samples from NYUv2 (https://cs.nyu.edu/~fergus/datasets/nyu_depth_v2.html)
+    NOTE: GT surface normals are from GeoNet (CVPR 2018) - https://github.com/xjqi/GeoNet
+"""
+import os
+import cv2
+import numpy as np
+
+from infer.dataset_normal import Sample
+
+
+def get_sample(base_data_dir, sample_path, info):
+    # e.g. sample_path = "test/000000_img.png"
+    scene_name = sample_path.split('/')[0]
+    img_name, img_ext = sample_path.split('/')[1].split('_img')
+
+    dataset_path = os.path.join(base_data_dir, 'dsine_eval', 'nyuv2')
+    img_path = '%s/%s' % (dataset_path, sample_path)
+    normal_png_path = img_path.replace('_img'+img_ext, '_normal.png')
+    normal_npy_path = img_path.replace('_img'+img_ext, '_normal.npy')
+    intrins_path = img_path.replace('_img'+img_ext, '_intrins.npy')
+    assert os.path.exists(img_path)
+
+    # read image (H, W, 3)
+    img = cv2.cvtColor(cv2.imread(img_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32) / 255.0
+    #保存图像
+    # cv2.imwrite(os.path.join(base_data_dir, img_name+'_img.png'), img*255)
+    
+    # read normal (H, W, 3)
+    if os.path.exists(normal_png_path):
+        normal = cv2.cvtColor(cv2.imread(normal_png_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+        normal_mask = np.sum(normal, axis=2, keepdims=True) > 0
+        normal = (normal.astype(np.float32) / 255.0) * 2.0 - 1.0
+    elif os.path.exists(normal_npy_path):
+        normal = np.load(normal_npy_path).astype(np.float32)
+        assert normal.ndim == 3 and normal.shape[2] == 3, f"Unexpected normal shape: {normal.shape}"
+        # GeoNet npy normals use opposite x-axis convention for this evaluation codepath.
+        normal[:, :, 0] *= -1.0
+        normal_mask = np.linalg.norm(normal, axis=2, keepdims=True) > 1e-6
+    else:
+        raise FileNotFoundError(f"Missing NYUv2 normal file: {normal_png_path} or {normal_npy_path}")
+
+    # read intrins (3, 3)
+    if os.path.exists(intrins_path):
+        intrins = np.load(intrins_path)
+    else:
+        # Fallback to NYUv2 default intrinsics used by many benchmarks.
+        intrins = np.array([
+            [518.8579, 0.0, 325.5824],
+            [0.0, 519.4696, 253.7362],
+            [0.0, 0.0, 1.0],
+        ], dtype=np.float32)
+
+    sample = Sample(
+        img=img,
+        normal=normal,
+        normal_mask=normal_mask,
+        intrins=intrins,
+
+        dataset_name='nyuv2',
+        scene_name=scene_name,
+        img_name=img_name,
+        info=info
+    )
+
+    return sample

FE2E/infer/dataset_normal/nyuv2/split/test.txt

@@ -0,0 +1,654 @@
+test/000000_img.png
+test/000001_img.png
+test/000008_img.png
+test/000013_img.png
+test/000014_img.png
+test/000015_img.png
+test/000016_img.png
+test/000017_img.png
+test/000020_img.png
+test/000027_img.png
+test/000028_img.png
+test/000029_img.png
+test/000030_img.png
+test/000031_img.png
+test/000032_img.png
+test/000033_img.png
+test/000034_img.png
+test/000035_img.png
+test/000036_img.png
+test/000037_img.png
+test/000038_img.png
+test/000039_img.png
+test/000040_img.png
+test/000041_img.png
+test/000042_img.png
+test/000045_img.png
+test/000046_img.png
+test/000055_img.png
+test/000056_img.png
+test/000058_img.png
+test/000059_img.png
+test/000060_img.png
+test/000061_img.png
+test/000062_img.png
+test/000075_img.png
+test/000076_img.png
+test/000077_img.png
+test/000078_img.png
+test/000083_img.png
+test/000084_img.png
+test/000085_img.png
+test/000086_img.png
+test/000087_img.png
+test/000088_img.png
+test/000089_img.png
+test/000090_img.png
+test/000116_img.png
+test/000117_img.png
+test/000118_img.png
+test/000124_img.png
+test/000125_img.png
+test/000126_img.png
+test/000127_img.png
+test/000128_img.png
+test/000130_img.png
+test/000131_img.png
+test/000132_img.png
+test/000133_img.png
+test/000136_img.png
+test/000152_img.png
+test/000153_img.png
+test/000154_img.png
+test/000166_img.png
+test/000167_img.png
+test/000168_img.png
+test/000170_img.png
+test/000171_img.png
+test/000172_img.png
+test/000173_img.png
+test/000174_img.png
+test/000175_img.png
+test/000179_img.png
+test/000180_img.png
+test/000181_img.png
+test/000182_img.png
+test/000183_img.png
+test/000184_img.png
+test/000185_img.png
+test/000186_img.png
+test/000187_img.png
+test/000188_img.png
+test/000189_img.png
+test/000190_img.png
+test/000191_img.png
+test/000192_img.png
+test/000193_img.png
+test/000194_img.png
+test/000195_img.png
+test/000196_img.png
+test/000197_img.png
+test/000198_img.png
+test/000199_img.png
+test/000200_img.png
+test/000201_img.png
+test/000206_img.png
+test/000207_img.png
+test/000208_img.png
+test/000209_img.png
+test/000210_img.png
+test/000211_img.png
+test/000219_img.png
+test/000220_img.png
+test/000221_img.png
+test/000249_img.png
+test/000263_img.png
+test/000270_img.png
+test/000271_img.png
+test/000272_img.png
+test/000278_img.png
+test/000279_img.png
+test/000280_img.png
+test/000281_img.png
+test/000282_img.png
+test/000283_img.png
+test/000284_img.png
+test/000295_img.png
+test/000296_img.png
+test/000297_img.png
+test/000298_img.png
+test/000299_img.png
+test/000300_img.png
+test/000301_img.png
+test/000309_img.png
+test/000310_img.png
+test/000311_img.png
+test/000314_img.png
+test/000315_img.png
+test/000316_img.png
+test/000324_img.png
+test/000325_img.png
+test/000326_img.png
+test/000327_img.png
+test/000328_img.png
+test/000329_img.png
+test/000330_img.png
+test/000331_img.png
+test/000332_img.png
+test/000333_img.png
+test/000334_img.png
+test/000350_img.png
+test/000351_img.png
+test/000354_img.png
+test/000355_img.png
+test/000356_img.png
+test/000357_img.png
+test/000358_img.png
+test/000359_img.png
+test/000360_img.png
+test/000361_img.png
+test/000362_img.png
+test/000363_img.png
+test/000383_img.png
+test/000384_img.png
+test/000385_img.png
+test/000386_img.png
+test/000387_img.png
+test/000388_img.png
+test/000389_img.png
+test/000394_img.png
+test/000395_img.png
+test/000396_img.png
+test/000410_img.png
+test/000411_img.png
+test/000412_img.png
+test/000413_img.png
+test/000429_img.png
+test/000430_img.png
+test/000431_img.png
+test/000432_img.png
+test/000433_img.png
+test/000434_img.png
+test/000440_img.png
+test/000441_img.png
+test/000442_img.png
+test/000443_img.png
+test/000444_img.png
+test/000445_img.png
+test/000446_img.png
+test/000447_img.png
+test/000461_img.png
+test/000462_img.png
+test/000463_img.png
+test/000464_img.png
+test/000465_img.png
+test/000468_img.png
+test/000469_img.png
+test/000470_img.png
+test/000471_img.png
+test/000472_img.png
+test/000473_img.png
+test/000474_img.png
+test/000475_img.png
+test/000476_img.png
+test/000507_img.png
+test/000508_img.png
+test/000509_img.png
+test/000510_img.png
+test/000511_img.png
+test/000512_img.png
+test/000514_img.png
+test/000515_img.png
+test/000516_img.png
+test/000517_img.png
+test/000518_img.png
+test/000519_img.png
+test/000520_img.png
+test/000521_img.png
+test/000522_img.png
+test/000523_img.png
+test/000524_img.png
+test/000525_img.png
+test/000530_img.png
+test/000531_img.png
+test/000532_img.png
+test/000536_img.png
+test/000537_img.png
+test/000538_img.png
+test/000548_img.png
+test/000549_img.png
+test/000550_img.png
+test/000554_img.png
+test/000555_img.png
+test/000556_img.png
+test/000557_img.png
+test/000558_img.png
+test/000559_img.png
+test/000560_img.png
+test/000561_img.png
+test/000562_img.png
+test/000563_img.png
+test/000564_img.png
+test/000565_img.png
+test/000566_img.png
+test/000567_img.png
+test/000568_img.png
+test/000569_img.png
+test/000570_img.png
+test/000578_img.png
+test/000579_img.png
+test/000580_img.png
+test/000581_img.png
+test/000582_img.png
+test/000590_img.png
+test/000591_img.png
+test/000592_img.png
+test/000593_img.png
+test/000602_img.png
+test/000603_img.png
+test/000604_img.png
+test/000605_img.png
+test/000606_img.png
+test/000611_img.png
+test/000612_img.png
+test/000616_img.png
+test/000617_img.png
+test/000618_img.png
+test/000619_img.png
+test/000620_img.png
+test/000632_img.png
+test/000633_img.png
+test/000634_img.png
+test/000635_img.png
+test/000636_img.png
+test/000637_img.png
+test/000643_img.png
+test/000644_img.png
+test/000649_img.png
+test/000650_img.png
+test/000655_img.png
+test/000656_img.png
+test/000657_img.png
+test/000662_img.png
+test/000663_img.png
+test/000667_img.png
+test/000668_img.png
+test/000669_img.png
+test/000670_img.png
+test/000671_img.png
+test/000672_img.png
+test/000675_img.png
+test/000676_img.png
+test/000677_img.png
+test/000678_img.png
+test/000679_img.png
+test/000680_img.png
+test/000685_img.png
+test/000686_img.png
+test/000687_img.png
+test/000688_img.png
+test/000689_img.png
+test/000692_img.png
+test/000693_img.png
+test/000696_img.png
+test/000697_img.png
+test/000698_img.png
+test/000705_img.png
+test/000706_img.png
+test/000707_img.png
+test/000708_img.png
+test/000709_img.png
+test/000710_img.png
+test/000711_img.png
+test/000712_img.png
+test/000716_img.png
+test/000717_img.png
+test/000723_img.png
+test/000724_img.png
+test/000725_img.png
+test/000726_img.png
+test/000727_img.png
+test/000730_img.png
+test/000731_img.png
+test/000732_img.png
+test/000733_img.png
+test/000742_img.png
+test/000743_img.png
+test/000758_img.png
+test/000759_img.png
+test/000760_img.png
+test/000761_img.png
+test/000762_img.png
+test/000763_img.png
+test/000764_img.png
+test/000765_img.png
+test/000766_img.png
+test/000767_img.png
+test/000768_img.png
+test/000769_img.png
+test/000770_img.png
+test/000771_img.png
+test/000772_img.png
+test/000773_img.png
+test/000774_img.png
+test/000775_img.png
+test/000776_img.png
+test/000777_img.png
+test/000778_img.png
+test/000779_img.png
+test/000780_img.png
+test/000781_img.png
+test/000782_img.png
+test/000783_img.png
+test/000784_img.png
+test/000785_img.png
+test/000786_img.png
+test/000799_img.png
+test/000800_img.png
+test/000801_img.png
+test/000802_img.png
+test/000803_img.png
+test/000809_img.png
+test/000810_img.png
+test/000811_img.png
+test/000812_img.png
+test/000813_img.png
+test/000820_img.png
+test/000821_img.png
+test/000822_img.png
+test/000832_img.png
+test/000833_img.png
+test/000834_img.png
+test/000835_img.png
+test/000836_img.png
+test/000837_img.png
+test/000838_img.png
+test/000839_img.png
+test/000840_img.png
+test/000841_img.png
+test/000842_img.png
+test/000843_img.png
+test/000844_img.png
+test/000845_img.png
+test/000849_img.png
+test/000850_img.png
+test/000851_img.png
+test/000856_img.png
+test/000857_img.png
+test/000858_img.png
+test/000859_img.png
+test/000860_img.png
+test/000861_img.png
+test/000868_img.png
+test/000869_img.png
+test/000870_img.png
+test/000905_img.png
+test/000906_img.png
+test/000907_img.png
+test/000916_img.png
+test/000917_img.png
+test/000918_img.png
+test/000925_img.png
+test/000926_img.png
+test/000927_img.png
+test/000931_img.png
+test/000932_img.png
+test/000933_img.png
+test/000934_img.png
+test/000944_img.png
+test/000945_img.png
+test/000946_img.png
+test/000958_img.png
+test/000959_img.png
+test/000960_img.png
+test/000961_img.png
+test/000964_img.png
+test/000965_img.png
+test/000966_img.png
+test/000969_img.png
+test/000970_img.png
+test/000971_img.png
+test/000972_img.png
+test/000973_img.png
+test/000974_img.png
+test/000975_img.png
+test/000976_img.png
+test/000990_img.png
+test/000991_img.png
+test/000992_img.png
+test/000993_img.png
+test/000994_img.png
+test/001000_img.png
+test/001001_img.png
+test/001002_img.png
+test/001003_img.png
+test/001009_img.png
+test/001010_img.png
+test/001011_img.png
+test/001020_img.png
+test/001021_img.png
+test/001022_img.png
+test/001031_img.png
+test/001032_img.png
+test/001033_img.png
+test/001037_img.png
+test/001038_img.png
+test/001047_img.png
+test/001048_img.png
+test/001051_img.png
+test/001052_img.png
+test/001056_img.png
+test/001057_img.png
+test/001074_img.png
+test/001075_img.png
+test/001076_img.png
+test/001077_img.png
+test/001078_img.png
+test/001079_img.png
+test/001080_img.png
+test/001081_img.png
+test/001082_img.png
+test/001083_img.png
+test/001087_img.png
+test/001088_img.png
+test/001089_img.png
+test/001090_img.png
+test/001091_img.png
+test/001092_img.png
+test/001093_img.png
+test/001094_img.png
+test/001095_img.png
+test/001097_img.png
+test/001098_img.png
+test/001099_img.png
+test/001100_img.png
+test/001101_img.png
+test/001102_img.png
+test/001103_img.png
+test/001105_img.png
+test/001106_img.png
+test/001107_img.png
+test/001108_img.png
+test/001116_img.png
+test/001117_img.png
+test/001118_img.png
+test/001122_img.png
+test/001123_img.png
+test/001124_img.png
+test/001125_img.png
+test/001126_img.png
+test/001127_img.png
+test/001128_img.png
+test/001129_img.png
+test/001130_img.png
+test/001134_img.png
+test/001135_img.png
+test/001143_img.png
+test/001144_img.png
+test/001145_img.png
+test/001146_img.png
+test/001147_img.png
+test/001148_img.png
+test/001149_img.png
+test/001150_img.png
+test/001151_img.png
+test/001152_img.png
+test/001153_img.png
+test/001154_img.png
+test/001155_img.png
+test/001156_img.png
+test/001157_img.png
+test/001161_img.png
+test/001162_img.png
+test/001163_img.png
+test/001164_img.png
+test/001165_img.png
+test/001166_img.png
+test/001169_img.png
+test/001170_img.png
+test/001173_img.png
+test/001174_img.png
+test/001175_img.png
+test/001178_img.png
+test/001179_img.png
+test/001180_img.png
+test/001181_img.png
+test/001182_img.png
+test/001183_img.png
+test/001191_img.png
+test/001192_img.png
+test/001193_img.png
+test/001194_img.png
+test/001195_img.png
+test/001200_img.png
+test/001201_img.png
+test/001202_img.png
+test/001203_img.png
+test/001204_img.png
+test/001205_img.png
+test/001206_img.png
+test/001207_img.png
+test/001208_img.png
+test/001209_img.png
+test/001210_img.png
+test/001211_img.png
+test/001215_img.png
+test/001216_img.png
+test/001217_img.png
+test/001218_img.png
+test/001219_img.png
+test/001225_img.png
+test/001226_img.png
+test/001227_img.png
+test/001228_img.png
+test/001229_img.png
+test/001232_img.png
+test/001233_img.png
+test/001234_img.png
+test/001246_img.png
+test/001247_img.png
+test/001248_img.png
+test/001249_img.png
+test/001253_img.png
+test/001254_img.png
+test/001255_img.png
+test/001256_img.png
+test/001257_img.png
+test/001258_img.png
+test/001259_img.png
+test/001260_img.png
+test/001261_img.png
+test/001262_img.png
+test/001263_img.png
+test/001264_img.png
+test/001274_img.png
+test/001275_img.png
+test/001276_img.png
+test/001277_img.png
+test/001278_img.png
+test/001279_img.png
+test/001284_img.png
+test/001285_img.png
+test/001286_img.png
+test/001287_img.png
+test/001288_img.png
+test/001289_img.png
+test/001290_img.png
+test/001291_img.png
+test/001292_img.png
+test/001293_img.png
+test/001294_img.png
+test/001296_img.png
+test/001297_img.png
+test/001298_img.png
+test/001301_img.png
+test/001302_img.png
+test/001303_img.png
+test/001304_img.png
+test/001305_img.png
+test/001306_img.png
+test/001307_img.png
+test/001313_img.png
+test/001314_img.png
+test/001328_img.png
+test/001329_img.png
+test/001330_img.png
+test/001331_img.png
+test/001334_img.png
+test/001335_img.png
+test/001336_img.png
+test/001337_img.png
+test/001338_img.png
+test/001339_img.png
+test/001346_img.png
+test/001347_img.png
+test/001348_img.png
+test/001352_img.png
+test/001353_img.png
+test/001354_img.png
+test/001355_img.png
+test/001363_img.png
+test/001364_img.png
+test/001367_img.png
+test/001368_img.png
+test/001383_img.png
+test/001384_img.png
+test/001385_img.png
+test/001386_img.png
+test/001387_img.png
+test/001388_img.png
+test/001389_img.png
+test/001390_img.png
+test/001393_img.png
+test/001394_img.png
+test/001395_img.png
+test/001396_img.png
+test/001397_img.png
+test/001398_img.png
+test/001399_img.png
+test/001400_img.png
+test/001406_img.png
+test/001407_img.png
+test/001408_img.png
+test/001409_img.png
+test/001410_img.png
+test/001411_img.png
+test/001412_img.png
+test/001413_img.png
+test/001420_img.png
+test/001421_img.png
+test/001422_img.png
+test/001423_img.png
+test/001429_img.png
+test/001430_img.png
+test/001431_img.png
+test/001432_img.png
+test/001440_img.png
+test/001441_img.png
+test/001442_img.png
+test/001443_img.png
+test/001444_img.png
+test/001445_img.png
+test/001446_img.png
+test/001447_img.png
+test/001448_img.png

FE2E/infer/dataset_normal/nyuv2/split/train.txt

@@ -0,0 +1,795 @@
+train/000002_img.png
+train/000003_img.png
+train/000004_img.png
+train/000005_img.png
+train/000006_img.png
+train/000007_img.png
+train/000009_img.png
+train/000010_img.png
+train/000011_img.png
+train/000012_img.png
+train/000018_img.png
+train/000019_img.png
+train/000021_img.png
+train/000022_img.png
+train/000023_img.png
+train/000024_img.png
+train/000025_img.png
+train/000026_img.png
+train/000043_img.png
+train/000044_img.png
+train/000047_img.png
+train/000048_img.png
+train/000049_img.png
+train/000050_img.png
+train/000051_img.png
+train/000052_img.png
+train/000053_img.png
+train/000054_img.png
+train/000057_img.png
+train/000063_img.png
+train/000064_img.png
+train/000065_img.png
+train/000066_img.png
+train/000067_img.png
+train/000068_img.png
+train/000069_img.png
+train/000070_img.png
+train/000071_img.png
+train/000072_img.png
+train/000073_img.png
+train/000074_img.png
+train/000079_img.png
+train/000080_img.png
+train/000081_img.png
+train/000082_img.png
+train/000091_img.png
+train/000092_img.png
+train/000093_img.png
+train/000094_img.png
+train/000095_img.png
+train/000096_img.png
+train/000097_img.png
+train/000098_img.png
+train/000099_img.png
+train/000100_img.png
+train/000101_img.png
+train/000102_img.png
+train/000103_img.png
+train/000104_img.png
+train/000105_img.png
+train/000106_img.png
+train/000107_img.png
+train/000108_img.png
+train/000109_img.png
+train/000110_img.png
+train/000111_img.png
+train/000112_img.png
+train/000113_img.png
+train/000114_img.png
+train/000115_img.png
+train/000119_img.png
+train/000120_img.png
+train/000121_img.png
+train/000122_img.png
+train/000123_img.png
+train/000129_img.png
+train/000134_img.png
+train/000135_img.png
+train/000137_img.png
+train/000138_img.png
+train/000139_img.png
+train/000140_img.png
+train/000141_img.png
+train/000142_img.png
+train/000143_img.png
+train/000144_img.png
+train/000145_img.png
+train/000146_img.png
+train/000147_img.png
+train/000148_img.png
+train/000149_img.png
+train/000150_img.png
+train/000151_img.png
+train/000155_img.png
+train/000156_img.png
+train/000157_img.png
+train/000158_img.png
+train/000159_img.png
+train/000160_img.png
+train/000161_img.png
+train/000162_img.png
+train/000163_img.png
+train/000164_img.png
+train/000165_img.png
+train/000169_img.png
+train/000176_img.png
+train/000177_img.png
+train/000178_img.png
+train/000202_img.png
+train/000203_img.png
+train/000204_img.png
+train/000205_img.png
+train/000212_img.png
+train/000213_img.png
+train/000214_img.png
+train/000215_img.png
+train/000216_img.png
+train/000217_img.png
+train/000218_img.png
+train/000222_img.png
+train/000223_img.png
+train/000224_img.png
+train/000225_img.png
+train/000226_img.png
+train/000227_img.png
+train/000228_img.png
+train/000229_img.png
+train/000230_img.png
+train/000231_img.png
+train/000232_img.png
+train/000233_img.png
+train/000234_img.png
+train/000235_img.png
+train/000236_img.png
+train/000237_img.png
+train/000238_img.png
+train/000239_img.png
+train/000240_img.png
+train/000241_img.png
+train/000242_img.png
+train/000243_img.png
+train/000244_img.png
+train/000245_img.png
+train/000246_img.png
+train/000247_img.png
+train/000248_img.png
+train/000250_img.png
+train/000251_img.png
+train/000252_img.png
+train/000253_img.png
+train/000254_img.png
+train/000255_img.png
+train/000256_img.png
+train/000257_img.png
+train/000258_img.png
+train/000259_img.png
+train/000260_img.png
+train/000261_img.png
+train/000262_img.png
+train/000264_img.png
+train/000265_img.png
+train/000266_img.png
+train/000267_img.png
+train/000268_img.png
+train/000269_img.png
+train/000273_img.png
+train/000274_img.png
+train/000275_img.png
+train/000276_img.png
+train/000277_img.png
+train/000285_img.png
+train/000286_img.png
+train/000287_img.png
+train/000288_img.png
+train/000289_img.png
+train/000290_img.png
+train/000291_img.png
+train/000292_img.png
+train/000293_img.png
+train/000294_img.png
+train/000302_img.png
+train/000303_img.png
+train/000304_img.png
+train/000305_img.png
+train/000306_img.png
+train/000307_img.png
+train/000308_img.png
+train/000312_img.png
+train/000313_img.png
+train/000317_img.png
+train/000318_img.png
+train/000319_img.png
+train/000320_img.png
+train/000321_img.png
+train/000322_img.png
+train/000323_img.png
+train/000335_img.png
+train/000336_img.png
+train/000337_img.png
+train/000338_img.png
+train/000339_img.png
+train/000340_img.png
+train/000341_img.png
+train/000342_img.png
+train/000343_img.png
+train/000344_img.png
+train/000345_img.png
+train/000346_img.png
+train/000347_img.png
+train/000348_img.png
+train/000349_img.png
+train/000352_img.png
+train/000353_img.png
+train/000364_img.png
+train/000365_img.png
+train/000366_img.png
+train/000367_img.png
+train/000368_img.png
+train/000369_img.png
+train/000370_img.png
+train/000371_img.png
+train/000372_img.png
+train/000373_img.png
+train/000374_img.png
+train/000375_img.png
+train/000376_img.png
+train/000377_img.png
+train/000378_img.png
+train/000379_img.png
+train/000380_img.png
+train/000381_img.png
+train/000382_img.png
+train/000390_img.png
+train/000391_img.png
+train/000392_img.png
+train/000393_img.png
+train/000397_img.png
+train/000398_img.png
+train/000399_img.png
+train/000400_img.png
+train/000401_img.png
+train/000402_img.png
+train/000403_img.png
+train/000404_img.png
+train/000405_img.png
+train/000406_img.png
+train/000407_img.png
+train/000408_img.png
+train/000409_img.png
+train/000414_img.png
+train/000415_img.png
+train/000416_img.png
+train/000417_img.png
+train/000418_img.png
+train/000419_img.png
+train/000420_img.png
+train/000421_img.png
+train/000422_img.png
+train/000423_img.png
+train/000424_img.png
+train/000425_img.png
+train/000426_img.png
+train/000427_img.png
+train/000428_img.png
+train/000435_img.png
+train/000436_img.png
+train/000437_img.png
+train/000438_img.png
+train/000439_img.png
+train/000448_img.png
+train/000449_img.png
+train/000450_img.png
+train/000451_img.png
+train/000452_img.png
+train/000453_img.png
+train/000454_img.png
+train/000455_img.png
+train/000456_img.png
+train/000457_img.png
+train/000458_img.png
+train/000459_img.png
+train/000460_img.png
+train/000466_img.png
+train/000467_img.png
+train/000477_img.png
+train/000478_img.png
+train/000479_img.png
+train/000480_img.png
+train/000481_img.png
+train/000482_img.png
+train/000483_img.png
+train/000484_img.png
+train/000485_img.png
+train/000486_img.png
+train/000487_img.png
+train/000488_img.png
+train/000489_img.png
+train/000490_img.png
+train/000491_img.png
+train/000492_img.png
+train/000493_img.png
+train/000494_img.png
+train/000495_img.png
+train/000496_img.png
+train/000497_img.png
+train/000498_img.png
+train/000499_img.png
+train/000500_img.png
+train/000501_img.png
+train/000502_img.png
+train/000503_img.png
+train/000504_img.png
+train/000505_img.png
+train/000506_img.png
+train/000513_img.png
+train/000526_img.png
+train/000527_img.png
+train/000528_img.png
+train/000529_img.png
+train/000533_img.png
+train/000534_img.png
+train/000535_img.png
+train/000539_img.png
+train/000540_img.png
+train/000541_img.png
+train/000542_img.png
+train/000543_img.png
+train/000544_img.png
+train/000545_img.png
+train/000546_img.png
+train/000547_img.png
+train/000551_img.png
+train/000552_img.png
+train/000553_img.png
+train/000571_img.png
+train/000572_img.png
+train/000573_img.png
+train/000574_img.png
+train/000575_img.png
+train/000576_img.png
+train/000577_img.png
+train/000583_img.png
+train/000584_img.png
+train/000585_img.png
+train/000586_img.png
+train/000587_img.png
+train/000588_img.png
+train/000589_img.png
+train/000594_img.png
+train/000595_img.png
+train/000596_img.png
+train/000597_img.png
+train/000598_img.png
+train/000599_img.png
+train/000600_img.png
+train/000601_img.png
+train/000607_img.png
+train/000608_img.png
+train/000609_img.png
+train/000610_img.png
+train/000613_img.png
+train/000614_img.png
+train/000615_img.png
+train/000621_img.png
+train/000622_img.png
+train/000623_img.png
+train/000624_img.png
+train/000625_img.png
+train/000626_img.png
+train/000627_img.png
+train/000628_img.png
+train/000629_img.png
+train/000630_img.png
+train/000631_img.png
+train/000638_img.png
+train/000639_img.png
+train/000640_img.png
+train/000641_img.png
+train/000642_img.png
+train/000645_img.png
+train/000646_img.png
+train/000647_img.png
+train/000648_img.png
+train/000651_img.png
+train/000652_img.png
+train/000653_img.png
+train/000654_img.png
+train/000658_img.png
+train/000659_img.png
+train/000660_img.png
+train/000661_img.png
+train/000664_img.png
+train/000665_img.png
+train/000666_img.png
+train/000673_img.png
+train/000674_img.png
+train/000681_img.png
+train/000682_img.png
+train/000683_img.png
+train/000684_img.png
+train/000690_img.png
+train/000691_img.png
+train/000694_img.png
+train/000695_img.png
+train/000699_img.png
+train/000700_img.png
+train/000701_img.png
+train/000702_img.png
+train/000703_img.png
+train/000704_img.png
+train/000713_img.png
+train/000714_img.png
+train/000715_img.png
+train/000718_img.png
+train/000719_img.png
+train/000720_img.png
+train/000721_img.png
+train/000722_img.png
+train/000728_img.png
+train/000729_img.png
+train/000734_img.png
+train/000735_img.png
+train/000736_img.png
+train/000737_img.png
+train/000738_img.png
+train/000739_img.png
+train/000740_img.png
+train/000741_img.png
+train/000744_img.png
+train/000745_img.png
+train/000746_img.png
+train/000747_img.png
+train/000748_img.png
+train/000749_img.png
+train/000750_img.png
+train/000751_img.png
+train/000752_img.png
+train/000753_img.png
+train/000754_img.png
+train/000755_img.png
+train/000756_img.png
+train/000757_img.png
+train/000787_img.png
+train/000788_img.png
+train/000789_img.png
+train/000790_img.png
+train/000791_img.png
+train/000792_img.png
+train/000793_img.png
+train/000794_img.png
+train/000795_img.png
+train/000796_img.png
+train/000797_img.png
+train/000798_img.png
+train/000804_img.png
+train/000805_img.png
+train/000806_img.png
+train/000807_img.png
+train/000808_img.png
+train/000814_img.png
+train/000815_img.png
+train/000816_img.png
+train/000817_img.png
+train/000818_img.png
+train/000819_img.png
+train/000823_img.png
+train/000824_img.png
+train/000825_img.png
+train/000826_img.png
+train/000827_img.png
+train/000828_img.png
+train/000829_img.png
+train/000830_img.png
+train/000831_img.png
+train/000846_img.png
+train/000847_img.png
+train/000848_img.png
+train/000852_img.png
+train/000853_img.png
+train/000854_img.png
+train/000855_img.png
+train/000862_img.png
+train/000863_img.png
+train/000864_img.png
+train/000865_img.png
+train/000866_img.png
+train/000867_img.png
+train/000871_img.png
+train/000872_img.png
+train/000873_img.png
+train/000874_img.png
+train/000875_img.png
+train/000876_img.png
+train/000877_img.png
+train/000878_img.png
+train/000879_img.png
+train/000880_img.png
+train/000881_img.png
+train/000882_img.png
+train/000883_img.png
+train/000884_img.png
+train/000885_img.png
+train/000886_img.png
+train/000887_img.png
+train/000888_img.png
+train/000889_img.png
+train/000890_img.png
+train/000891_img.png
+train/000892_img.png
+train/000893_img.png
+train/000894_img.png
+train/000895_img.png
+train/000896_img.png
+train/000897_img.png
+train/000898_img.png
+train/000899_img.png
+train/000900_img.png
+train/000901_img.png
+train/000902_img.png
+train/000903_img.png
+train/000904_img.png
+train/000908_img.png
+train/000909_img.png
+train/000910_img.png
+train/000911_img.png
+train/000912_img.png
+train/000913_img.png
+train/000914_img.png
+train/000915_img.png
+train/000919_img.png
+train/000920_img.png
+train/000921_img.png
+train/000922_img.png
+train/000923_img.png
+train/000924_img.png
+train/000928_img.png
+train/000929_img.png
+train/000930_img.png
+train/000935_img.png
+train/000936_img.png
+train/000937_img.png
+train/000938_img.png
+train/000939_img.png
+train/000940_img.png
+train/000941_img.png
+train/000942_img.png
+train/000943_img.png
+train/000947_img.png
+train/000948_img.png
+train/000949_img.png
+train/000950_img.png
+train/000951_img.png
+train/000952_img.png
+train/000953_img.png
+train/000954_img.png
+train/000955_img.png
+train/000956_img.png
+train/000957_img.png
+train/000962_img.png
+train/000963_img.png
+train/000967_img.png
+train/000968_img.png
+train/000977_img.png
+train/000978_img.png
+train/000979_img.png
+train/000980_img.png
+train/000981_img.png
+train/000982_img.png
+train/000983_img.png
+train/000984_img.png
+train/000985_img.png
+train/000986_img.png
+train/000987_img.png
+train/000988_img.png
+train/000989_img.png
+train/000995_img.png
+train/000996_img.png
+train/000997_img.png
+train/000998_img.png
+train/000999_img.png
+train/001004_img.png
+train/001005_img.png
+train/001006_img.png
+train/001007_img.png
+train/001008_img.png
+train/001012_img.png
+train/001013_img.png
+train/001014_img.png
+train/001015_img.png
+train/001016_img.png
+train/001017_img.png
+train/001018_img.png
+train/001019_img.png
+train/001023_img.png
+train/001024_img.png
+train/001025_img.png
+train/001026_img.png
+train/001027_img.png
+train/001028_img.png
+train/001029_img.png
+train/001030_img.png
+train/001034_img.png
+train/001035_img.png
+train/001036_img.png
+train/001039_img.png
+train/001040_img.png
+train/001041_img.png
+train/001042_img.png
+train/001043_img.png
+train/001044_img.png
+train/001045_img.png
+train/001046_img.png
+train/001049_img.png
+train/001050_img.png
+train/001053_img.png
+train/001054_img.png
+train/001055_img.png
+train/001058_img.png
+train/001059_img.png
+train/001060_img.png
+train/001061_img.png
+train/001062_img.png
+train/001063_img.png
+train/001064_img.png
+train/001065_img.png
+train/001066_img.png
+train/001067_img.png
+train/001068_img.png
+train/001069_img.png
+train/001070_img.png
+train/001071_img.png
+train/001072_img.png
+train/001073_img.png
+train/001084_img.png
+train/001085_img.png
+train/001086_img.png
+train/001096_img.png
+train/001104_img.png
+train/001109_img.png
+train/001110_img.png
+train/001111_img.png
+train/001112_img.png
+train/001113_img.png
+train/001114_img.png
+train/001115_img.png
+train/001119_img.png
+train/001120_img.png
+train/001121_img.png
+train/001131_img.png
+train/001132_img.png
+train/001133_img.png
+train/001136_img.png
+train/001137_img.png
+train/001138_img.png
+train/001139_img.png
+train/001140_img.png
+train/001141_img.png
+train/001142_img.png
+train/001158_img.png
+train/001159_img.png
+train/001160_img.png
+train/001167_img.png
+train/001168_img.png
+train/001171_img.png
+train/001172_img.png
+train/001176_img.png
+train/001177_img.png
+train/001184_img.png
+train/001185_img.png
+train/001186_img.png
+train/001187_img.png
+train/001188_img.png
+train/001189_img.png
+train/001190_img.png
+train/001196_img.png
+train/001197_img.png
+train/001198_img.png
+train/001199_img.png
+train/001212_img.png
+train/001213_img.png
+train/001214_img.png
+train/001220_img.png
+train/001221_img.png
+train/001222_img.png
+train/001223_img.png
+train/001224_img.png
+train/001230_img.png
+train/001231_img.png
+train/001235_img.png
+train/001236_img.png
+train/001237_img.png
+train/001238_img.png
+train/001239_img.png
+train/001240_img.png
+train/001241_img.png
+train/001242_img.png
+train/001243_img.png
+train/001244_img.png
+train/001245_img.png
+train/001250_img.png
+train/001251_img.png
+train/001252_img.png
+train/001265_img.png
+train/001266_img.png
+train/001267_img.png
+train/001268_img.png
+train/001269_img.png
+train/001270_img.png
+train/001271_img.png
+train/001272_img.png
+train/001273_img.png
+train/001280_img.png
+train/001281_img.png
+train/001282_img.png
+train/001283_img.png
+train/001295_img.png
+train/001299_img.png
+train/001300_img.png
+train/001308_img.png
+train/001309_img.png
+train/001310_img.png
+train/001311_img.png
+train/001312_img.png
+train/001315_img.png
+train/001316_img.png
+train/001317_img.png
+train/001318_img.png
+train/001319_img.png
+train/001320_img.png
+train/001321_img.png
+train/001322_img.png
+train/001323_img.png
+train/001324_img.png
+train/001325_img.png
+train/001326_img.png
+train/001327_img.png
+train/001332_img.png
+train/001333_img.png
+train/001340_img.png
+train/001341_img.png
+train/001342_img.png
+train/001343_img.png
+train/001344_img.png
+train/001345_img.png
+train/001349_img.png
+train/001350_img.png
+train/001351_img.png
+train/001356_img.png
+train/001357_img.png
+train/001358_img.png
+train/001359_img.png
+train/001360_img.png
+train/001361_img.png
+train/001362_img.png
+train/001365_img.png
+train/001366_img.png
+train/001369_img.png
+train/001370_img.png
+train/001371_img.png
+train/001372_img.png
+train/001373_img.png
+train/001374_img.png
+train/001375_img.png
+train/001376_img.png
+train/001377_img.png
+train/001378_img.png
+train/001379_img.png
+train/001380_img.png
+train/001381_img.png
+train/001382_img.png
+train/001391_img.png
+train/001392_img.png
+train/001401_img.png
+train/001402_img.png
+train/001403_img.png
+train/001404_img.png
+train/001405_img.png
+train/001414_img.png
+train/001415_img.png
+train/001416_img.png
+train/001417_img.png
+train/001418_img.png
+train/001419_img.png
+train/001424_img.png
+train/001425_img.png
+train/001426_img.png
+train/001427_img.png
+train/001428_img.png
+train/001433_img.png
+train/001434_img.png
+train/001435_img.png
+train/001436_img.png
+train/001437_img.png
+train/001438_img.png
+train/001439_img.png

FE2E/infer/dataset_normal/oasis/__init__.py

@@ -0,0 +1,76 @@
+""" Get samples from OASIS validation set (https://pvl.cs.princeton.edu/OASIS/)
+"""
+import os
+import cv2
+import numpy as np
+import pickle
+
+from infer.dataset_normal import Sample
+
+
+def read_normal(path, h, w):
+    normal_dict = pickle.load(open(path, 'rb'))
+
+    mask = np.zeros((h,w))
+    normal = np.zeros((h,w,3))
+
+    # Stuff ROI normal into bounding box
+    min_y = normal_dict['min_y']
+    max_y = normal_dict['max_y']
+    min_x = normal_dict['min_x']
+    max_x = normal_dict['max_x']
+    roi_normal = normal_dict['normal']
+
+    # to LUB
+    normal[min_y:max_y+1, min_x:max_x+1, :] = roi_normal
+    normal = normal.astype(np.float32)
+    normal[:,:,0] *= -1
+    normal[:,:,1] *= -1
+
+    # Make mask
+    roi_mask = np.logical_or(np.logical_or(roi_normal[:,:,0] != 0, roi_normal[:,:,1] != 0), roi_normal[:,:,2] != 0).astype(np.float32)
+    mask[min_y:max_y+1, min_x:max_x+1] = roi_mask
+    mask = mask[:, :, None]
+    mask = mask > 0.5
+
+    return normal, mask
+
+
+def get_sample(base_data_dir, sample_path, info):
+    # e.g. sample_path = "val/100277_DT_img.png"
+    scene_name = sample_path.split('/')[0]
+    img_name, img_ext = sample_path.split('/')[-1].split('_img')
+
+    dataset_path = os.path.join(base_data_dir, 'dsine_eval', 'oasis')
+    img_path = '%s/%s' % (dataset_path, sample_path)
+    normal_path = img_path.replace('_img'+img_ext, '_normal.pkl')
+    intrins_path = img_path.replace('_img'+img_ext, '_intrins.npy')
+    assert os.path.exists(img_path)
+    assert os.path.exists(normal_path)
+    assert os.path.exists(intrins_path)
+
+    # read image (H, W, 3)
+    img = cv2.cvtColor(cv2.imread(img_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32) / 255.0
+
+    # read normal (H, W, 3)
+    h = img.shape[0]
+    w = img.shape[1]
+    normal, normal_mask = read_normal(normal_path, h, w)
+
+    # read intrins (3, 3)
+    intrins = np.load(intrins_path)
+
+    sample = Sample(
+        img=img,
+        normal=normal,
+        normal_mask=normal_mask,
+        intrins=intrins,
+
+        dataset_name='oasis',
+        scene_name=scene_name,
+        img_name=img_name,
+        info=info
+    )
+
+    return sample

FE2E/infer/dataset_normal/scannet/__init__.py

@@ -0,0 +1,64 @@
+""" Get samples from ScanNet (https://github.com/ScanNet/ScanNet)
+    NOTE: GT surface normals and data split are from FrameNet (ICCV 2019) - https://github.com/hjwdzh/FrameNet
+"""
+import os
+import cv2
+import numpy as np
+
+from infer.dataset_normal import Sample
+
+
+def get_sample(base_data_dir, sample_path, info):
+    # e.g. sample_path = "scene0532_00/000000_img.png"
+    scene_name = sample_path.split('/')[0]
+    img_name, img_ext = sample_path.split('/')[1].split('_img')
+
+    dataset_path = os.path.join(base_data_dir, 'dsine_eval', 'scannet')
+    img_path = '%s/%s' % (dataset_path, sample_path)
+    normal_png_path = img_path.replace('_img'+img_ext, '_normal.png')
+    normal_npy_path = img_path.replace('_img'+img_ext, '_normal.npy')
+    intrins_path = img_path.replace('_img'+img_ext, '_intrins.npy')
+    assert os.path.exists(img_path)
+
+    # read image (H, W, 3)
+    img = cv2.cvtColor(cv2.imread(img_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32) / 255.0
+
+    # read normal (H, W, 3)
+    if os.path.exists(normal_png_path):
+        normal = cv2.cvtColor(cv2.imread(normal_png_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+        normal_mask = np.sum(normal, axis=2, keepdims=True) > 0
+        normal = (normal.astype(np.float32) / 255.0) * 2.0 - 1.0
+    elif os.path.exists(normal_npy_path):
+        normal = np.load(normal_npy_path).astype(np.float32)
+        assert normal.ndim == 3 and normal.shape[2] == 3, f"Unexpected normal shape: {normal.shape}"
+        # FrameNet npy normals use opposite x-axis convention for this evaluation codepath.
+        normal[:, :, 0] *= -1.0
+        normal_mask = np.linalg.norm(normal, axis=2, keepdims=True) > 1e-6
+    else:
+        raise FileNotFoundError(f"Missing ScanNet normal file: {normal_png_path} or {normal_npy_path}")
+
+    # read intrins (3, 3)
+    if os.path.exists(intrins_path):
+        intrins = np.load(intrins_path)
+    else:
+        # Fallback intrinsics for ScanNet benchmark-sized frames.
+        intrins = np.array([
+            [577.870605, 0.0, 319.5],
+            [0.0, 577.870605, 239.5],
+            [0.0, 0.0, 1.0],
+        ], dtype=np.float32)
+
+    sample = Sample(
+        img=img,
+        normal=normal,
+        normal_mask=normal_mask,
+        intrins=intrins,
+
+        dataset_name='scannet',
+        scene_name=scene_name,
+        img_name=img_name,
+        info=info
+    )
+
+    return sample

FE2E/infer/dataset_normal/scannet/split/test.txt

@@ -0,0 +1,300 @@
+scene0532_00/000000_img.png
+scene0220_02/000450_img.png
+scene0061_00/000420_img.png
+scene0656_03/000490_img.png
+scene0525_01/000960_img.png
+scene0310_01/000630_img.png
+scene0628_00/000100_img.png
+scene0481_01/000900_img.png
+scene0608_02/001690_img.png
+scene0415_00/000750_img.png
+scene0603_00/001590_img.png
+scene0740_00/000510_img.png
+scene0740_00/002510_img.png
+scene0177_00/000550_img.png
+scene0395_01/000160_img.png
+scene0395_01/002160_img.png
+scene0343_00/000210_img.png
+scene0241_00/000970_img.png
+scene0717_00/000010_img.png
+scene0362_00/000890_img.png
+scene0457_02/000070_img.png
+scene0296_00/000040_img.png
+scene0157_00/000300_img.png
+scene0502_02/001130_img.png
+scene0101_01/000410_img.png
+scene0610_02/000730_img.png
+scene0610_02/002780_img.png
+scene0416_00/000950_img.png
+scene0685_02/000840_img.png
+scene0685_02/003050_img.png
+scene0662_01/001810_img.png
+scene0670_01/000840_img.png
+scene0670_01/002840_img.png
+scene0012_01/000160_img.png
+scene0012_01/002160_img.png
+scene0320_03/000790_img.png
+scene0780_00/000450_img.png
+scene0575_01/000510_img.png
+scene0575_01/002510_img.png
+scene0168_01/000100_img.png
+scene0335_02/000980_img.png
+scene0691_00/000400_img.png
+scene0778_00/001230_img.png
+scene0019_01/000490_img.png
+scene0554_00/001960_img.png
+scene0114_00/000260_img.png
+scene0457_01/000350_img.png
+scene0548_01/000150_img.png
+scene0548_01/002150_img.png
+scene0378_00/000610_img.png
+scene0505_00/000710_img.png
+scene0505_00/002710_img.png
+scene0695_01/001200_img.png
+scene0424_02/001060_img.png
+scene0627_00/001470_img.png
+scene0279_00/001200_img.png
+scene0291_01/001280_img.png
+scene0768_00/000330_img.png
+scene0768_00/002330_img.png
+scene0696_01/000340_img.png
+scene0145_00/000420_img.png
+scene0185_00/000370_img.png
+scene0185_00/002370_img.png
+scene0806_00/001070_img.png
+scene0261_03/000410_img.png
+scene0068_00/000570_img.png
+scene0804_00/000510_img.png
+scene0221_01/000490_img.png
+scene0630_03/000880_img.png
+scene0617_00/001270_img.png
+scene0206_02/001640_img.png
+scene0582_02/000590_img.png
+scene0031_01/000100_img.png
+scene0031_01/002100_img.png
+scene0505_04/001100_img.png
+scene0092_01/000600_img.png
+scene0542_00/000570_img.png
+scene0525_02/000480_img.png
+scene0756_00/000870_img.png
+scene0756_00/002870_img.png
+scene0061_01/001360_img.png
+scene0588_01/000940_img.png
+scene0256_02/000610_img.png
+scene0300_00/001430_img.png
+scene0451_05/001110_img.png
+scene0397_01/001390_img.png
+scene0372_00/000250_img.png
+scene0372_00/002280_img.png
+scene0647_01/000790_img.png
+scene0474_01/001600_img.png
+scene0404_00/000360_img.png
+scene0404_00/002360_img.png
+scene0702_01/000010_img.png
+scene0114_02/001420_img.png
+scene0460_00/000820_img.png
+scene0151_01/000390_img.png
+scene0151_01/002390_img.png
+scene0078_02/000520_img.png
+scene0772_00/001530_img.png
+scene0705_01/000490_img.png
+scene0408_01/000590_img.png
+scene0538_00/001430_img.png
+scene0651_00/000080_img.png
+scene0210_01/001040_img.png
+scene0002_00/000240_img.png
+scene0002_00/002240_img.png
+scene0002_00/004240_img.png
+scene0347_00/001040_img.png
+scene0072_02/000210_img.png
+scene0371_00/000410_img.png
+scene0211_00/001120_img.png
+scene0514_01/000050_img.png
+scene0473_01/000910_img.png
+scene0362_02/001970_img.png
+scene0024_02/001180_img.png
+scene0673_00/000280_img.png
+scene0176_00/000030_img.png
+scene0073_00/000400_img.png
+scene0186_01/000230_img.png
+scene0236_00/000490_img.png
+scene0210_00/000780_img.png
+scene0325_01/000920_img.png
+scene0044_01/000170_img.png
+scene0197_00/001300_img.png
+scene0472_02/001190_img.png
+scene0611_01/000010_img.png
+scene0758_00/000710_img.png
+scene0650_00/000870_img.png
+scene0728_00/000160_img.png
+scene0217_00/000450_img.png
+scene0335_01/000070_img.png
+scene0335_01/002070_img.png
+scene0504_00/001230_img.png
+scene0669_01/000780_img.png
+scene0143_01/000760_img.png
+scene0320_01/000460_img.png
+scene0645_01/000340_img.png
+scene0645_01/002340_img.png
+scene0645_01/004340_img.png
+scene0248_01/000300_img.png
+scene0416_04/000580_img.png
+scene0416_04/002580_img.png
+scene0534_00/000780_img.png
+scene0489_02/001010_img.png
+scene0379_00/000500_img.png
+scene0500_01/000480_img.png
+scene0561_01/000850_img.png
+scene0580_00/001200_img.png
+scene0580_00/003200_img.png
+scene0552_01/000430_img.png
+scene0744_00/000020_img.png
+scene0744_00/002020_img.png
+scene0044_00/001030_img.png
+scene0301_02/000140_img.png
+scene0422_00/000940_img.png
+scene0655_01/000700_img.png
+scene0803_00/001770_img.png
+scene0629_02/000530_img.png
+scene0609_01/000030_img.png
+scene0623_01/000930_img.png
+scene0596_00/000660_img.png
+scene0150_01/000150_img.png
+scene0501_02/001770_img.png
+scene0560_00/000420_img.png
+scene0406_02/000220_img.png
+scene0587_02/001010_img.png
+scene0540_00/000730_img.png
+scene0074_01/001390_img.png
+scene0620_01/000200_img.png
+scene0486_00/000320_img.png
+scene0486_00/002320_img.png
+scene0600_02/001350_img.png
+scene0085_00/001800_img.png
+scene0084_00/000560_img.png
+scene0319_00/000580_img.png
+scene0400_01/000890_img.png
+scene0548_02/000020_img.png
+scene0548_02/002110_img.png
+scene0244_01/000360_img.png
+scene0785_00/000630_img.png
+scene0785_00/002630_img.png
+scene0479_02/000530_img.png
+scene0121_00/000310_img.png
+scene0107_00/000410_img.png
+scene0328_00/000040_img.png
+scene0196_00/000350_img.png
+scene0404_01/000720_img.png
+scene0404_01/002720_img.png
+scene0081_02/000710_img.png
+scene0666_00/000360_img.png
+scene0367_00/000330_img.png
+scene0340_01/000460_img.png
+scene0300_01/000980_img.png
+scene0275_00/001230_img.png
+scene0036_00/000440_img.png
+scene0520_01/001350_img.png
+scene0113_01/000000_img.png
+scene0541_01/000020_img.png
+scene0034_01/001070_img.png
+scene0030_01/000160_img.png
+scene0438_00/000510_img.png
+scene0679_00/001210_img.png
+scene0546_00/000400_img.png
+scene0223_00/000600_img.png
+scene0403_00/001570_img.png
+scene0001_01/001170_img.png
+scene0014_00/001880_img.png
+scene0673_02/001480_img.png
+scene0794_00/000000_img.png
+scene0209_02/000210_img.png
+scene0801_00/000180_img.png
+scene0086_00/000800_img.png
+scene0501_00/001370_img.png
+scene0412_01/000070_img.png
+scene0339_00/000730_img.png
+scene0724_00/000500_img.png
+scene0654_01/001080_img.png
+scene0081_01/000460_img.png
+scene0576_02/001030_img.png
+scene0589_01/001010_img.png
+scene0428_00/001470_img.png
+scene0199_00/000110_img.png
+scene0513_00/000440_img.png
+scene0512_00/000110_img.png
+scene0508_00/000840_img.png
+scene0009_02/000360_img.png
+scene0303_00/001250_img.png
+scene0533_01/001240_img.png
+scene0445_01/000640_img.png
+scene0392_00/001870_img.png
+scene0111_01/001560_img.png
+scene0192_02/000620_img.png
+scene0396_00/000170_img.png
+scene0376_02/000320_img.png
+scene0100_01/000880_img.png
+scene0578_01/000920_img.png
+scene0765_00/000400_img.png
+scene0784_00/000420_img.png
+scene0784_00/002430_img.png
+scene0784_00/004430_img.png
+scene0065_02/000740_img.png
+scene0207_02/000040_img.png
+scene0207_02/002040_img.png
+scene0440_01/000500_img.png
+scene0220_01/000450_img.png
+scene0713_00/000000_img.png
+scene0713_00/002000_img.png
+scene0697_02/000070_img.png
+scene0697_02/002070_img.png
+scene0286_03/001170_img.png
+scene0393_02/000810_img.png
+scene0370_02/000080_img.png
+scene0370_02/002080_img.png
+scene0452_02/000950_img.png
+scene0226_00/001270_img.png
+scene0474_03/001610_img.png
+scene0304_00/000600_img.png
+scene0250_01/000820_img.png
+scene0250_01/002820_img.png
+scene0419_02/001480_img.png
+scene0548_00/000000_img.png
+scene0548_00/002050_img.png
+scene0568_00/000150_img.png
+scene0324_00/000080_img.png
+scene0567_00/000070_img.png
+scene0567_00/002070_img.png
+scene0659_01/001460_img.png
+scene0540_02/000120_img.png
+scene0488_01/000680_img.png
+scene0146_01/000260_img.png
+scene0580_01/001130_img.png
+scene0580_01/003130_img.png
+scene0421_02/001320_img.png
+scene0657_00/000090_img.png
+scene0588_02/001380_img.png
+scene0471_02/000570_img.png
+scene0186_00/000260_img.png
+scene0186_00/002260_img.png
+scene0065_01/000800_img.png
+scene0456_01/000710_img.png
+scene0524_01/001680_img.png
+scene0062_00/000690_img.png
+scene0134_01/000780_img.png
+scene0294_00/000280_img.png
+scene0294_00/002280_img.png
+scene0452_01/001090_img.png
+scene0142_00/000770_img.png
+scene0135_00/000330_img.png
+scene0279_02/000780_img.png
+scene0698_01/000390_img.png
+scene0520_00/000820_img.png
+scene0448_01/000850_img.png
+scene0276_00/000320_img.png
+scene0731_00/000220_img.png
+scene0599_01/001090_img.png
+scene0631_02/000210_img.png
+scene0748_00/000890_img.png
+scene0170_02/000190_img.png
+scene0215_01/000400_img.png
+scene0178_00/001330_img.png

FE2E/infer/dataset_normal/sintel/__init__.py

@@ -0,0 +1,49 @@
+""" Get samples from Sintel (http://sintel.is.tue.mpg.de/)
+    NOTE: We computed the GT surface normals by doing discontinuity-aware plane fitting
+"""
+import os
+import cv2
+import numpy as np
+os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
+
+from infer.dataset_normal import Sample
+
+
+
+def get_sample(base_data_dir, sample_path, info):
+    # e.g. sample_path = "alley_1/frame_0001_img.png"
+    scene_name = sample_path.split('/')[0]
+    img_name, img_ext = sample_path.split('/')[1].split('_img')
+
+    dataset_path = os.path.join(base_data_dir, 'dsine_eval', 'sintel')
+    img_path = '%s/%s' % (dataset_path, sample_path)
+    normal_path = img_path.replace('_img'+img_ext, '_normal.exr')
+    intrins_path = img_path.replace('_img'+img_ext, '_intrins.npy')
+    assert os.path.exists(img_path)
+    assert os.path.exists(normal_path)
+    assert os.path.exists(intrins_path)
+
+    # read image (H, W, 3)
+    img = cv2.cvtColor(cv2.imread(img_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32) / 255.0
+
+    # read normal (H, W, 3)
+    normal = cv2.cvtColor(cv2.imread(normal_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    normal_mask = np.linalg.norm(normal, axis=2, keepdims=True) > 0.5
+
+    # read intrins (3, 3)
+    intrins = np.load(intrins_path)
+
+    sample = Sample(
+        img=img,
+        normal=normal,
+        normal_mask=normal_mask,
+        intrins=intrins,
+
+        dataset_name='sintel',
+        scene_name=scene_name,
+        img_name=img_name,
+        info=info
+    )
+
+    return sample

FE2E/infer/dataset_normal/sintel/split/sintel.txt

@@ -0,0 +1,1064 @@
+alley_1/frame_0001_img.png
+alley_1/frame_0002_img.png
+alley_1/frame_0003_img.png
+alley_1/frame_0004_img.png
+alley_1/frame_0005_img.png
+alley_1/frame_0006_img.png
+alley_1/frame_0007_img.png
+alley_1/frame_0008_img.png
+alley_1/frame_0009_img.png
+alley_1/frame_0010_img.png
+alley_1/frame_0011_img.png
+alley_1/frame_0012_img.png
+alley_1/frame_0013_img.png
+alley_1/frame_0014_img.png
+alley_1/frame_0015_img.png
+alley_1/frame_0016_img.png
+alley_1/frame_0017_img.png
+alley_1/frame_0018_img.png
+alley_1/frame_0019_img.png
+alley_1/frame_0020_img.png
+alley_1/frame_0021_img.png
+alley_1/frame_0022_img.png
+alley_1/frame_0023_img.png
+alley_1/frame_0024_img.png
+alley_1/frame_0025_img.png
+alley_1/frame_0026_img.png
+alley_1/frame_0027_img.png
+alley_1/frame_0028_img.png
+alley_1/frame_0029_img.png
+alley_1/frame_0030_img.png
+alley_1/frame_0031_img.png
+alley_1/frame_0032_img.png
+alley_1/frame_0033_img.png
+alley_1/frame_0034_img.png
+alley_1/frame_0035_img.png
+alley_1/frame_0036_img.png
+alley_1/frame_0037_img.png
+alley_1/frame_0038_img.png
+alley_1/frame_0039_img.png
+alley_1/frame_0040_img.png
+alley_1/frame_0041_img.png
+alley_1/frame_0042_img.png
+alley_1/frame_0043_img.png
+alley_1/frame_0044_img.png
+alley_1/frame_0045_img.png
+alley_1/frame_0046_img.png
+alley_1/frame_0047_img.png
+alley_1/frame_0048_img.png
+alley_1/frame_0049_img.png
+alley_1/frame_0050_img.png
+alley_2/frame_0001_img.png
+alley_2/frame_0002_img.png
+alley_2/frame_0003_img.png
+alley_2/frame_0004_img.png
+alley_2/frame_0005_img.png
+alley_2/frame_0006_img.png
+alley_2/frame_0007_img.png
+alley_2/frame_0008_img.png
+alley_2/frame_0009_img.png
+alley_2/frame_0010_img.png
+alley_2/frame_0011_img.png
+alley_2/frame_0012_img.png
+alley_2/frame_0013_img.png
+alley_2/frame_0014_img.png
+alley_2/frame_0015_img.png
+alley_2/frame_0016_img.png
+alley_2/frame_0017_img.png
+alley_2/frame_0018_img.png
+alley_2/frame_0019_img.png
+alley_2/frame_0020_img.png
+alley_2/frame_0021_img.png
+alley_2/frame_0022_img.png
+alley_2/frame_0023_img.png
+alley_2/frame_0024_img.png
+alley_2/frame_0025_img.png
+alley_2/frame_0026_img.png
+alley_2/frame_0027_img.png
+alley_2/frame_0028_img.png
+alley_2/frame_0029_img.png
+alley_2/frame_0030_img.png
+alley_2/frame_0031_img.png
+alley_2/frame_0032_img.png
+alley_2/frame_0033_img.png
+alley_2/frame_0034_img.png
+alley_2/frame_0035_img.png
+alley_2/frame_0036_img.png
+alley_2/frame_0037_img.png
+alley_2/frame_0038_img.png
+alley_2/frame_0039_img.png
+alley_2/frame_0040_img.png
+alley_2/frame_0041_img.png
+alley_2/frame_0042_img.png
+alley_2/frame_0043_img.png
+alley_2/frame_0044_img.png
+alley_2/frame_0045_img.png
+alley_2/frame_0046_img.png
+alley_2/frame_0047_img.png
+alley_2/frame_0048_img.png
+alley_2/frame_0049_img.png
+alley_2/frame_0050_img.png
+ambush_2/frame_0001_img.png
+ambush_2/frame_0002_img.png
+ambush_2/frame_0003_img.png
+ambush_2/frame_0004_img.png
+ambush_2/frame_0005_img.png
+ambush_2/frame_0006_img.png
+ambush_2/frame_0007_img.png
+ambush_2/frame_0008_img.png
+ambush_2/frame_0009_img.png
+ambush_2/frame_0010_img.png
+ambush_2/frame_0011_img.png
+ambush_2/frame_0012_img.png
+ambush_2/frame_0013_img.png
+ambush_2/frame_0014_img.png
+ambush_2/frame_0015_img.png
+ambush_2/frame_0016_img.png
+ambush_2/frame_0017_img.png
+ambush_2/frame_0018_img.png
+ambush_2/frame_0019_img.png
+ambush_2/frame_0020_img.png
+ambush_2/frame_0021_img.png
+ambush_4/frame_0001_img.png
+ambush_4/frame_0002_img.png
+ambush_4/frame_0003_img.png
+ambush_4/frame_0004_img.png
+ambush_4/frame_0005_img.png
+ambush_4/frame_0006_img.png
+ambush_4/frame_0007_img.png
+ambush_4/frame_0008_img.png
+ambush_4/frame_0009_img.png
+ambush_4/frame_0010_img.png
+ambush_4/frame_0011_img.png
+ambush_4/frame_0012_img.png
+ambush_4/frame_0013_img.png
+ambush_4/frame_0014_img.png
+ambush_4/frame_0015_img.png
+ambush_4/frame_0016_img.png
+ambush_4/frame_0017_img.png
+ambush_4/frame_0018_img.png
+ambush_4/frame_0019_img.png
+ambush_4/frame_0020_img.png
+ambush_4/frame_0021_img.png
+ambush_4/frame_0022_img.png
+ambush_4/frame_0023_img.png
+ambush_4/frame_0024_img.png
+ambush_4/frame_0025_img.png
+ambush_4/frame_0026_img.png
+ambush_4/frame_0027_img.png
+ambush_4/frame_0028_img.png
+ambush_4/frame_0029_img.png
+ambush_4/frame_0030_img.png
+ambush_4/frame_0031_img.png
+ambush_4/frame_0032_img.png
+ambush_4/frame_0033_img.png
+ambush_5/frame_0001_img.png
+ambush_5/frame_0002_img.png
+ambush_5/frame_0003_img.png
+ambush_5/frame_0004_img.png
+ambush_5/frame_0005_img.png
+ambush_5/frame_0006_img.png
+ambush_5/frame_0007_img.png
+ambush_5/frame_0008_img.png
+ambush_5/frame_0009_img.png
+ambush_5/frame_0010_img.png
+ambush_5/frame_0011_img.png
+ambush_5/frame_0012_img.png
+ambush_5/frame_0013_img.png
+ambush_5/frame_0014_img.png
+ambush_5/frame_0015_img.png
+ambush_5/frame_0016_img.png
+ambush_5/frame_0017_img.png
+ambush_5/frame_0018_img.png
+ambush_5/frame_0019_img.png
+ambush_5/frame_0020_img.png
+ambush_5/frame_0021_img.png
+ambush_5/frame_0022_img.png
+ambush_5/frame_0023_img.png
+ambush_5/frame_0024_img.png
+ambush_5/frame_0025_img.png
+ambush_5/frame_0026_img.png
+ambush_5/frame_0027_img.png
+ambush_5/frame_0028_img.png
+ambush_5/frame_0029_img.png
+ambush_5/frame_0030_img.png
+ambush_5/frame_0031_img.png
+ambush_5/frame_0032_img.png
+ambush_5/frame_0033_img.png
+ambush_5/frame_0034_img.png
+ambush_5/frame_0035_img.png
+ambush_5/frame_0036_img.png
+ambush_5/frame_0037_img.png
+ambush_5/frame_0038_img.png
+ambush_5/frame_0039_img.png
+ambush_5/frame_0040_img.png
+ambush_5/frame_0041_img.png
+ambush_5/frame_0042_img.png
+ambush_5/frame_0043_img.png
+ambush_5/frame_0044_img.png
+ambush_5/frame_0045_img.png
+ambush_5/frame_0046_img.png
+ambush_5/frame_0047_img.png
+ambush_5/frame_0048_img.png
+ambush_5/frame_0049_img.png
+ambush_5/frame_0050_img.png
+ambush_6/frame_0001_img.png
+ambush_6/frame_0002_img.png
+ambush_6/frame_0003_img.png
+ambush_6/frame_0004_img.png
+ambush_6/frame_0005_img.png
+ambush_6/frame_0006_img.png
+ambush_6/frame_0007_img.png
+ambush_6/frame_0008_img.png
+ambush_6/frame_0009_img.png
+ambush_6/frame_0010_img.png
+ambush_6/frame_0011_img.png
+ambush_6/frame_0012_img.png
+ambush_6/frame_0013_img.png
+ambush_6/frame_0014_img.png
+ambush_6/frame_0015_img.png
+ambush_6/frame_0016_img.png
+ambush_6/frame_0017_img.png
+ambush_6/frame_0018_img.png
+ambush_6/frame_0019_img.png
+ambush_6/frame_0020_img.png
+ambush_7/frame_0001_img.png
+ambush_7/frame_0002_img.png
+ambush_7/frame_0003_img.png
+ambush_7/frame_0004_img.png
+ambush_7/frame_0005_img.png
+ambush_7/frame_0006_img.png
+ambush_7/frame_0007_img.png
+ambush_7/frame_0008_img.png
+ambush_7/frame_0009_img.png
+ambush_7/frame_0010_img.png
+ambush_7/frame_0011_img.png
+ambush_7/frame_0012_img.png
+ambush_7/frame_0013_img.png
+ambush_7/frame_0014_img.png
+ambush_7/frame_0015_img.png
+ambush_7/frame_0016_img.png
+ambush_7/frame_0017_img.png
+ambush_7/frame_0018_img.png
+ambush_7/frame_0019_img.png
+ambush_7/frame_0020_img.png
+ambush_7/frame_0021_img.png
+ambush_7/frame_0022_img.png
+ambush_7/frame_0023_img.png
+ambush_7/frame_0024_img.png
+ambush_7/frame_0025_img.png
+ambush_7/frame_0026_img.png
+ambush_7/frame_0027_img.png
+ambush_7/frame_0028_img.png
+ambush_7/frame_0029_img.png
+ambush_7/frame_0030_img.png
+ambush_7/frame_0031_img.png
+ambush_7/frame_0032_img.png
+ambush_7/frame_0033_img.png
+ambush_7/frame_0034_img.png
+ambush_7/frame_0035_img.png
+ambush_7/frame_0036_img.png
+ambush_7/frame_0037_img.png
+ambush_7/frame_0038_img.png
+ambush_7/frame_0039_img.png
+ambush_7/frame_0040_img.png
+ambush_7/frame_0041_img.png
+ambush_7/frame_0042_img.png
+ambush_7/frame_0043_img.png
+ambush_7/frame_0044_img.png
+ambush_7/frame_0045_img.png
+ambush_7/frame_0046_img.png
+ambush_7/frame_0047_img.png
+ambush_7/frame_0048_img.png
+ambush_7/frame_0049_img.png
+ambush_7/frame_0050_img.png
+bamboo_1/frame_0001_img.png
+bamboo_1/frame_0002_img.png
+bamboo_1/frame_0003_img.png
+bamboo_1/frame_0004_img.png
+bamboo_1/frame_0005_img.png
+bamboo_1/frame_0006_img.png
+bamboo_1/frame_0007_img.png
+bamboo_1/frame_0008_img.png
+bamboo_1/frame_0009_img.png
+bamboo_1/frame_0010_img.png
+bamboo_1/frame_0011_img.png
+bamboo_1/frame_0012_img.png
+bamboo_1/frame_0013_img.png
+bamboo_1/frame_0014_img.png
+bamboo_1/frame_0015_img.png
+bamboo_1/frame_0016_img.png
+bamboo_1/frame_0017_img.png
+bamboo_1/frame_0018_img.png
+bamboo_1/frame_0019_img.png
+bamboo_1/frame_0020_img.png
+bamboo_1/frame_0021_img.png
+bamboo_1/frame_0022_img.png
+bamboo_1/frame_0023_img.png
+bamboo_1/frame_0024_img.png
+bamboo_1/frame_0025_img.png
+bamboo_1/frame_0026_img.png
+bamboo_1/frame_0027_img.png
+bamboo_1/frame_0028_img.png
+bamboo_1/frame_0029_img.png
+bamboo_1/frame_0030_img.png
+bamboo_1/frame_0031_img.png
+bamboo_1/frame_0032_img.png
+bamboo_1/frame_0033_img.png
+bamboo_1/frame_0034_img.png
+bamboo_1/frame_0035_img.png
+bamboo_1/frame_0036_img.png
+bamboo_1/frame_0037_img.png
+bamboo_1/frame_0038_img.png
+bamboo_1/frame_0039_img.png
+bamboo_1/frame_0040_img.png
+bamboo_1/frame_0041_img.png
+bamboo_1/frame_0042_img.png
+bamboo_1/frame_0043_img.png
+bamboo_1/frame_0044_img.png
+bamboo_1/frame_0045_img.png
+bamboo_1/frame_0046_img.png
+bamboo_1/frame_0047_img.png
+bamboo_1/frame_0048_img.png
+bamboo_1/frame_0049_img.png
+bamboo_1/frame_0050_img.png
+bamboo_2/frame_0001_img.png
+bamboo_2/frame_0002_img.png
+bamboo_2/frame_0003_img.png
+bamboo_2/frame_0004_img.png
+bamboo_2/frame_0005_img.png
+bamboo_2/frame_0006_img.png
+bamboo_2/frame_0007_img.png
+bamboo_2/frame_0008_img.png
+bamboo_2/frame_0009_img.png
+bamboo_2/frame_0010_img.png
+bamboo_2/frame_0011_img.png
+bamboo_2/frame_0012_img.png
+bamboo_2/frame_0013_img.png
+bamboo_2/frame_0014_img.png
+bamboo_2/frame_0015_img.png
+bamboo_2/frame_0016_img.png
+bamboo_2/frame_0017_img.png
+bamboo_2/frame_0018_img.png
+bamboo_2/frame_0019_img.png
+bamboo_2/frame_0020_img.png
+bamboo_2/frame_0021_img.png
+bamboo_2/frame_0022_img.png
+bamboo_2/frame_0023_img.png
+bamboo_2/frame_0024_img.png
+bamboo_2/frame_0025_img.png
+bamboo_2/frame_0026_img.png
+bamboo_2/frame_0027_img.png
+bamboo_2/frame_0028_img.png
+bamboo_2/frame_0029_img.png
+bamboo_2/frame_0030_img.png
+bamboo_2/frame_0031_img.png
+bamboo_2/frame_0032_img.png
+bamboo_2/frame_0033_img.png
+bamboo_2/frame_0034_img.png
+bamboo_2/frame_0035_img.png
+bamboo_2/frame_0036_img.png
+bamboo_2/frame_0037_img.png
+bamboo_2/frame_0038_img.png
+bamboo_2/frame_0039_img.png
+bamboo_2/frame_0040_img.png
+bamboo_2/frame_0041_img.png
+bamboo_2/frame_0042_img.png
+bamboo_2/frame_0043_img.png
+bamboo_2/frame_0044_img.png
+bamboo_2/frame_0045_img.png
+bamboo_2/frame_0046_img.png
+bamboo_2/frame_0047_img.png
+bamboo_2/frame_0048_img.png
+bamboo_2/frame_0049_img.png
+bamboo_2/frame_0050_img.png
+bandage_1/frame_0001_img.png
+bandage_1/frame_0002_img.png
+bandage_1/frame_0003_img.png
+bandage_1/frame_0004_img.png
+bandage_1/frame_0005_img.png
+bandage_1/frame_0006_img.png
+bandage_1/frame_0007_img.png
+bandage_1/frame_0008_img.png
+bandage_1/frame_0009_img.png
+bandage_1/frame_0010_img.png
+bandage_1/frame_0011_img.png
+bandage_1/frame_0012_img.png
+bandage_1/frame_0013_img.png
+bandage_1/frame_0014_img.png
+bandage_1/frame_0015_img.png
+bandage_1/frame_0016_img.png
+bandage_1/frame_0017_img.png
+bandage_1/frame_0018_img.png
+bandage_1/frame_0019_img.png
+bandage_1/frame_0020_img.png
+bandage_1/frame_0021_img.png
+bandage_1/frame_0022_img.png
+bandage_1/frame_0023_img.png
+bandage_1/frame_0024_img.png
+bandage_1/frame_0025_img.png
+bandage_1/frame_0026_img.png
+bandage_1/frame_0027_img.png
+bandage_1/frame_0028_img.png
+bandage_1/frame_0029_img.png
+bandage_1/frame_0030_img.png
+bandage_1/frame_0031_img.png
+bandage_1/frame_0032_img.png
+bandage_1/frame_0033_img.png
+bandage_1/frame_0034_img.png
+bandage_1/frame_0035_img.png
+bandage_1/frame_0036_img.png
+bandage_1/frame_0037_img.png
+bandage_1/frame_0038_img.png
+bandage_1/frame_0039_img.png
+bandage_1/frame_0040_img.png
+bandage_1/frame_0041_img.png
+bandage_1/frame_0042_img.png
+bandage_1/frame_0043_img.png
+bandage_1/frame_0044_img.png
+bandage_1/frame_0045_img.png
+bandage_1/frame_0046_img.png
+bandage_1/frame_0047_img.png
+bandage_1/frame_0048_img.png
+bandage_1/frame_0049_img.png
+bandage_1/frame_0050_img.png
+bandage_2/frame_0001_img.png
+bandage_2/frame_0002_img.png
+bandage_2/frame_0003_img.png
+bandage_2/frame_0004_img.png
+bandage_2/frame_0005_img.png
+bandage_2/frame_0006_img.png
+bandage_2/frame_0007_img.png
+bandage_2/frame_0008_img.png
+bandage_2/frame_0009_img.png
+bandage_2/frame_0010_img.png
+bandage_2/frame_0011_img.png
+bandage_2/frame_0012_img.png
+bandage_2/frame_0013_img.png
+bandage_2/frame_0014_img.png
+bandage_2/frame_0015_img.png
+bandage_2/frame_0016_img.png
+bandage_2/frame_0017_img.png
+bandage_2/frame_0018_img.png
+bandage_2/frame_0019_img.png
+bandage_2/frame_0020_img.png
+bandage_2/frame_0021_img.png
+bandage_2/frame_0022_img.png
+bandage_2/frame_0023_img.png
+bandage_2/frame_0024_img.png
+bandage_2/frame_0025_img.png
+bandage_2/frame_0026_img.png
+bandage_2/frame_0027_img.png
+bandage_2/frame_0028_img.png
+bandage_2/frame_0029_img.png
+bandage_2/frame_0030_img.png
+bandage_2/frame_0031_img.png
+bandage_2/frame_0032_img.png
+bandage_2/frame_0033_img.png
+bandage_2/frame_0034_img.png
+bandage_2/frame_0035_img.png
+bandage_2/frame_0036_img.png
+bandage_2/frame_0037_img.png
+bandage_2/frame_0038_img.png
+bandage_2/frame_0039_img.png
+bandage_2/frame_0040_img.png
+bandage_2/frame_0041_img.png
+bandage_2/frame_0042_img.png
+bandage_2/frame_0043_img.png
+bandage_2/frame_0044_img.png
+bandage_2/frame_0045_img.png
+bandage_2/frame_0046_img.png
+bandage_2/frame_0047_img.png
+bandage_2/frame_0048_img.png
+bandage_2/frame_0049_img.png
+bandage_2/frame_0050_img.png
+cave_2/frame_0001_img.png
+cave_2/frame_0002_img.png
+cave_2/frame_0003_img.png
+cave_2/frame_0004_img.png
+cave_2/frame_0005_img.png
+cave_2/frame_0006_img.png
+cave_2/frame_0007_img.png
+cave_2/frame_0008_img.png
+cave_2/frame_0009_img.png
+cave_2/frame_0010_img.png
+cave_2/frame_0011_img.png
+cave_2/frame_0012_img.png
+cave_2/frame_0013_img.png
+cave_2/frame_0014_img.png
+cave_2/frame_0015_img.png
+cave_2/frame_0016_img.png
+cave_2/frame_0017_img.png
+cave_2/frame_0018_img.png
+cave_2/frame_0019_img.png
+cave_2/frame_0020_img.png
+cave_2/frame_0021_img.png
+cave_2/frame_0022_img.png
+cave_2/frame_0023_img.png
+cave_2/frame_0024_img.png
+cave_2/frame_0025_img.png
+cave_2/frame_0026_img.png
+cave_2/frame_0027_img.png
+cave_2/frame_0028_img.png
+cave_2/frame_0029_img.png
+cave_2/frame_0030_img.png
+cave_2/frame_0031_img.png
+cave_2/frame_0032_img.png
+cave_2/frame_0033_img.png
+cave_2/frame_0034_img.png
+cave_2/frame_0035_img.png
+cave_2/frame_0036_img.png
+cave_2/frame_0037_img.png
+cave_2/frame_0038_img.png
+cave_2/frame_0039_img.png
+cave_2/frame_0040_img.png
+cave_2/frame_0041_img.png
+cave_2/frame_0042_img.png
+cave_2/frame_0043_img.png
+cave_2/frame_0044_img.png
+cave_2/frame_0045_img.png
+cave_2/frame_0046_img.png
+cave_2/frame_0047_img.png
+cave_2/frame_0048_img.png
+cave_2/frame_0049_img.png
+cave_2/frame_0050_img.png
+cave_4/frame_0001_img.png
+cave_4/frame_0002_img.png
+cave_4/frame_0003_img.png
+cave_4/frame_0004_img.png
+cave_4/frame_0005_img.png
+cave_4/frame_0006_img.png
+cave_4/frame_0007_img.png
+cave_4/frame_0008_img.png
+cave_4/frame_0009_img.png
+cave_4/frame_0010_img.png
+cave_4/frame_0011_img.png
+cave_4/frame_0012_img.png
+cave_4/frame_0013_img.png
+cave_4/frame_0014_img.png
+cave_4/frame_0015_img.png
+cave_4/frame_0016_img.png
+cave_4/frame_0017_img.png
+cave_4/frame_0018_img.png
+cave_4/frame_0019_img.png
+cave_4/frame_0020_img.png
+cave_4/frame_0021_img.png
+cave_4/frame_0022_img.png
+cave_4/frame_0023_img.png
+cave_4/frame_0024_img.png
+cave_4/frame_0025_img.png
+cave_4/frame_0026_img.png
+cave_4/frame_0027_img.png
+cave_4/frame_0028_img.png
+cave_4/frame_0029_img.png
+cave_4/frame_0030_img.png
+cave_4/frame_0031_img.png
+cave_4/frame_0032_img.png
+cave_4/frame_0033_img.png
+cave_4/frame_0034_img.png
+cave_4/frame_0035_img.png
+cave_4/frame_0036_img.png
+cave_4/frame_0037_img.png
+cave_4/frame_0038_img.png
+cave_4/frame_0039_img.png
+cave_4/frame_0040_img.png
+cave_4/frame_0041_img.png
+cave_4/frame_0042_img.png
+cave_4/frame_0043_img.png
+cave_4/frame_0044_img.png
+cave_4/frame_0045_img.png
+cave_4/frame_0046_img.png
+cave_4/frame_0047_img.png
+cave_4/frame_0048_img.png
+cave_4/frame_0049_img.png
+cave_4/frame_0050_img.png
+market_2/frame_0001_img.png
+market_2/frame_0002_img.png
+market_2/frame_0003_img.png
+market_2/frame_0004_img.png
+market_2/frame_0005_img.png
+market_2/frame_0006_img.png
+market_2/frame_0007_img.png
+market_2/frame_0008_img.png
+market_2/frame_0009_img.png
+market_2/frame_0010_img.png
+market_2/frame_0011_img.png
+market_2/frame_0012_img.png
+market_2/frame_0013_img.png
+market_2/frame_0014_img.png
+market_2/frame_0015_img.png
+market_2/frame_0016_img.png
+market_2/frame_0017_img.png
+market_2/frame_0018_img.png
+market_2/frame_0019_img.png
+market_2/frame_0020_img.png
+market_2/frame_0021_img.png
+market_2/frame_0022_img.png
+market_2/frame_0023_img.png
+market_2/frame_0024_img.png
+market_2/frame_0025_img.png
+market_2/frame_0026_img.png
+market_2/frame_0027_img.png
+market_2/frame_0028_img.png
+market_2/frame_0029_img.png
+market_2/frame_0030_img.png
+market_2/frame_0031_img.png
+market_2/frame_0032_img.png
+market_2/frame_0033_img.png
+market_2/frame_0034_img.png
+market_2/frame_0035_img.png
+market_2/frame_0036_img.png
+market_2/frame_0037_img.png
+market_2/frame_0038_img.png
+market_2/frame_0039_img.png
+market_2/frame_0040_img.png
+market_2/frame_0041_img.png
+market_2/frame_0042_img.png
+market_2/frame_0043_img.png
+market_2/frame_0044_img.png
+market_2/frame_0045_img.png
+market_2/frame_0046_img.png
+market_2/frame_0047_img.png
+market_2/frame_0048_img.png
+market_2/frame_0049_img.png
+market_2/frame_0050_img.png
+market_5/frame_0001_img.png
+market_5/frame_0002_img.png
+market_5/frame_0003_img.png
+market_5/frame_0004_img.png
+market_5/frame_0005_img.png
+market_5/frame_0006_img.png
+market_5/frame_0007_img.png
+market_5/frame_0008_img.png
+market_5/frame_0009_img.png
+market_5/frame_0010_img.png
+market_5/frame_0011_img.png
+market_5/frame_0012_img.png
+market_5/frame_0013_img.png
+market_5/frame_0014_img.png
+market_5/frame_0015_img.png
+market_5/frame_0016_img.png
+market_5/frame_0017_img.png
+market_5/frame_0018_img.png
+market_5/frame_0019_img.png
+market_5/frame_0020_img.png
+market_5/frame_0021_img.png
+market_5/frame_0022_img.png
+market_5/frame_0023_img.png
+market_5/frame_0024_img.png
+market_5/frame_0025_img.png
+market_5/frame_0026_img.png
+market_5/frame_0027_img.png
+market_5/frame_0028_img.png
+market_5/frame_0029_img.png
+market_5/frame_0030_img.png
+market_5/frame_0031_img.png
+market_5/frame_0032_img.png
+market_5/frame_0033_img.png
+market_5/frame_0034_img.png
+market_5/frame_0035_img.png
+market_5/frame_0036_img.png
+market_5/frame_0037_img.png
+market_5/frame_0038_img.png
+market_5/frame_0039_img.png
+market_5/frame_0040_img.png
+market_5/frame_0041_img.png
+market_5/frame_0042_img.png
+market_5/frame_0043_img.png
+market_5/frame_0044_img.png
+market_5/frame_0045_img.png
+market_5/frame_0046_img.png
+market_5/frame_0047_img.png
+market_5/frame_0048_img.png
+market_5/frame_0049_img.png
+market_5/frame_0050_img.png
+market_6/frame_0001_img.png
+market_6/frame_0002_img.png
+market_6/frame_0003_img.png
+market_6/frame_0004_img.png
+market_6/frame_0005_img.png
+market_6/frame_0006_img.png
+market_6/frame_0007_img.png
+market_6/frame_0008_img.png
+market_6/frame_0009_img.png
+market_6/frame_0010_img.png
+market_6/frame_0011_img.png
+market_6/frame_0012_img.png
+market_6/frame_0013_img.png
+market_6/frame_0014_img.png
+market_6/frame_0015_img.png
+market_6/frame_0016_img.png
+market_6/frame_0017_img.png
+market_6/frame_0018_img.png
+market_6/frame_0019_img.png
+market_6/frame_0020_img.png
+market_6/frame_0021_img.png
+market_6/frame_0022_img.png
+market_6/frame_0023_img.png
+market_6/frame_0024_img.png
+market_6/frame_0025_img.png
+market_6/frame_0026_img.png
+market_6/frame_0027_img.png
+market_6/frame_0028_img.png
+market_6/frame_0029_img.png
+market_6/frame_0030_img.png
+market_6/frame_0031_img.png
+market_6/frame_0032_img.png
+market_6/frame_0033_img.png
+market_6/frame_0034_img.png
+market_6/frame_0035_img.png
+market_6/frame_0036_img.png
+market_6/frame_0037_img.png
+market_6/frame_0038_img.png
+market_6/frame_0039_img.png
+market_6/frame_0040_img.png
+mountain_1/frame_0001_img.png
+mountain_1/frame_0002_img.png
+mountain_1/frame_0003_img.png
+mountain_1/frame_0004_img.png
+mountain_1/frame_0005_img.png
+mountain_1/frame_0006_img.png
+mountain_1/frame_0007_img.png
+mountain_1/frame_0008_img.png
+mountain_1/frame_0009_img.png
+mountain_1/frame_0010_img.png
+mountain_1/frame_0011_img.png
+mountain_1/frame_0012_img.png
+mountain_1/frame_0013_img.png
+mountain_1/frame_0014_img.png
+mountain_1/frame_0015_img.png
+mountain_1/frame_0016_img.png
+mountain_1/frame_0017_img.png
+mountain_1/frame_0018_img.png
+mountain_1/frame_0019_img.png
+mountain_1/frame_0020_img.png
+mountain_1/frame_0021_img.png
+mountain_1/frame_0022_img.png
+mountain_1/frame_0023_img.png
+mountain_1/frame_0024_img.png
+mountain_1/frame_0025_img.png
+mountain_1/frame_0026_img.png
+mountain_1/frame_0027_img.png
+mountain_1/frame_0028_img.png
+mountain_1/frame_0029_img.png
+mountain_1/frame_0030_img.png
+mountain_1/frame_0031_img.png
+mountain_1/frame_0032_img.png
+mountain_1/frame_0033_img.png
+mountain_1/frame_0034_img.png
+mountain_1/frame_0035_img.png
+mountain_1/frame_0036_img.png
+mountain_1/frame_0037_img.png
+mountain_1/frame_0038_img.png
+mountain_1/frame_0039_img.png
+mountain_1/frame_0040_img.png
+mountain_1/frame_0041_img.png
+mountain_1/frame_0042_img.png
+mountain_1/frame_0043_img.png
+mountain_1/frame_0044_img.png
+mountain_1/frame_0045_img.png
+mountain_1/frame_0046_img.png
+mountain_1/frame_0047_img.png
+mountain_1/frame_0048_img.png
+mountain_1/frame_0049_img.png
+mountain_1/frame_0050_img.png
+shaman_2/frame_0001_img.png
+shaman_2/frame_0002_img.png
+shaman_2/frame_0003_img.png
+shaman_2/frame_0004_img.png
+shaman_2/frame_0005_img.png
+shaman_2/frame_0006_img.png
+shaman_2/frame_0007_img.png
+shaman_2/frame_0008_img.png
+shaman_2/frame_0009_img.png
+shaman_2/frame_0010_img.png
+shaman_2/frame_0011_img.png
+shaman_2/frame_0012_img.png
+shaman_2/frame_0013_img.png
+shaman_2/frame_0014_img.png
+shaman_2/frame_0015_img.png
+shaman_2/frame_0016_img.png
+shaman_2/frame_0017_img.png
+shaman_2/frame_0018_img.png
+shaman_2/frame_0019_img.png
+shaman_2/frame_0020_img.png
+shaman_2/frame_0021_img.png
+shaman_2/frame_0022_img.png
+shaman_2/frame_0023_img.png
+shaman_2/frame_0024_img.png
+shaman_2/frame_0025_img.png
+shaman_2/frame_0026_img.png
+shaman_2/frame_0027_img.png
+shaman_2/frame_0028_img.png
+shaman_2/frame_0029_img.png
+shaman_2/frame_0030_img.png
+shaman_2/frame_0031_img.png
+shaman_2/frame_0032_img.png
+shaman_2/frame_0033_img.png
+shaman_2/frame_0034_img.png
+shaman_2/frame_0035_img.png
+shaman_2/frame_0036_img.png
+shaman_2/frame_0037_img.png
+shaman_2/frame_0038_img.png
+shaman_2/frame_0039_img.png
+shaman_2/frame_0040_img.png
+shaman_2/frame_0041_img.png
+shaman_2/frame_0042_img.png
+shaman_2/frame_0043_img.png
+shaman_2/frame_0044_img.png
+shaman_2/frame_0045_img.png
+shaman_2/frame_0046_img.png
+shaman_2/frame_0047_img.png
+shaman_2/frame_0048_img.png
+shaman_2/frame_0049_img.png
+shaman_2/frame_0050_img.png
+shaman_3/frame_0001_img.png
+shaman_3/frame_0002_img.png
+shaman_3/frame_0003_img.png
+shaman_3/frame_0004_img.png
+shaman_3/frame_0005_img.png
+shaman_3/frame_0006_img.png
+shaman_3/frame_0007_img.png
+shaman_3/frame_0008_img.png
+shaman_3/frame_0009_img.png
+shaman_3/frame_0010_img.png
+shaman_3/frame_0011_img.png
+shaman_3/frame_0012_img.png
+shaman_3/frame_0013_img.png
+shaman_3/frame_0014_img.png
+shaman_3/frame_0015_img.png
+shaman_3/frame_0016_img.png
+shaman_3/frame_0017_img.png
+shaman_3/frame_0018_img.png
+shaman_3/frame_0019_img.png
+shaman_3/frame_0020_img.png
+shaman_3/frame_0021_img.png
+shaman_3/frame_0022_img.png
+shaman_3/frame_0023_img.png
+shaman_3/frame_0024_img.png
+shaman_3/frame_0025_img.png
+shaman_3/frame_0026_img.png
+shaman_3/frame_0027_img.png
+shaman_3/frame_0028_img.png
+shaman_3/frame_0029_img.png
+shaman_3/frame_0030_img.png
+shaman_3/frame_0031_img.png
+shaman_3/frame_0032_img.png
+shaman_3/frame_0033_img.png
+shaman_3/frame_0034_img.png
+shaman_3/frame_0035_img.png
+shaman_3/frame_0036_img.png
+shaman_3/frame_0037_img.png
+shaman_3/frame_0038_img.png
+shaman_3/frame_0039_img.png
+shaman_3/frame_0040_img.png
+shaman_3/frame_0041_img.png
+shaman_3/frame_0042_img.png
+shaman_3/frame_0043_img.png
+shaman_3/frame_0044_img.png
+shaman_3/frame_0045_img.png
+shaman_3/frame_0046_img.png
+shaman_3/frame_0047_img.png
+shaman_3/frame_0048_img.png
+shaman_3/frame_0049_img.png
+shaman_3/frame_0050_img.png
+sleeping_1/frame_0001_img.png
+sleeping_1/frame_0002_img.png
+sleeping_1/frame_0003_img.png
+sleeping_1/frame_0004_img.png
+sleeping_1/frame_0005_img.png
+sleeping_1/frame_0006_img.png
+sleeping_1/frame_0007_img.png
+sleeping_1/frame_0008_img.png
+sleeping_1/frame_0009_img.png
+sleeping_1/frame_0010_img.png
+sleeping_1/frame_0011_img.png
+sleeping_1/frame_0012_img.png
+sleeping_1/frame_0013_img.png
+sleeping_1/frame_0014_img.png
+sleeping_1/frame_0015_img.png
+sleeping_1/frame_0016_img.png
+sleeping_1/frame_0017_img.png
+sleeping_1/frame_0018_img.png
+sleeping_1/frame_0019_img.png
+sleeping_1/frame_0020_img.png
+sleeping_1/frame_0021_img.png
+sleeping_1/frame_0022_img.png
+sleeping_1/frame_0023_img.png
+sleeping_1/frame_0024_img.png
+sleeping_1/frame_0025_img.png
+sleeping_1/frame_0026_img.png
+sleeping_1/frame_0027_img.png
+sleeping_1/frame_0028_img.png
+sleeping_1/frame_0029_img.png
+sleeping_1/frame_0030_img.png
+sleeping_1/frame_0031_img.png
+sleeping_1/frame_0032_img.png
+sleeping_1/frame_0033_img.png
+sleeping_1/frame_0034_img.png
+sleeping_1/frame_0035_img.png
+sleeping_1/frame_0036_img.png
+sleeping_1/frame_0037_img.png
+sleeping_1/frame_0038_img.png
+sleeping_1/frame_0039_img.png
+sleeping_1/frame_0040_img.png
+sleeping_1/frame_0041_img.png
+sleeping_1/frame_0042_img.png
+sleeping_1/frame_0043_img.png
+sleeping_1/frame_0044_img.png
+sleeping_1/frame_0045_img.png
+sleeping_1/frame_0046_img.png
+sleeping_1/frame_0047_img.png
+sleeping_1/frame_0048_img.png
+sleeping_1/frame_0049_img.png
+sleeping_1/frame_0050_img.png
+sleeping_2/frame_0001_img.png
+sleeping_2/frame_0002_img.png
+sleeping_2/frame_0003_img.png
+sleeping_2/frame_0004_img.png
+sleeping_2/frame_0005_img.png
+sleeping_2/frame_0006_img.png
+sleeping_2/frame_0007_img.png
+sleeping_2/frame_0008_img.png
+sleeping_2/frame_0009_img.png
+sleeping_2/frame_0010_img.png
+sleeping_2/frame_0011_img.png
+sleeping_2/frame_0012_img.png
+sleeping_2/frame_0013_img.png
+sleeping_2/frame_0014_img.png
+sleeping_2/frame_0015_img.png
+sleeping_2/frame_0016_img.png
+sleeping_2/frame_0017_img.png
+sleeping_2/frame_0018_img.png
+sleeping_2/frame_0019_img.png
+sleeping_2/frame_0020_img.png
+sleeping_2/frame_0021_img.png
+sleeping_2/frame_0022_img.png
+sleeping_2/frame_0023_img.png
+sleeping_2/frame_0024_img.png
+sleeping_2/frame_0025_img.png
+sleeping_2/frame_0026_img.png
+sleeping_2/frame_0027_img.png
+sleeping_2/frame_0028_img.png
+sleeping_2/frame_0029_img.png
+sleeping_2/frame_0030_img.png
+sleeping_2/frame_0031_img.png
+sleeping_2/frame_0032_img.png
+sleeping_2/frame_0033_img.png
+sleeping_2/frame_0034_img.png
+sleeping_2/frame_0035_img.png
+sleeping_2/frame_0036_img.png
+sleeping_2/frame_0037_img.png
+sleeping_2/frame_0038_img.png
+sleeping_2/frame_0039_img.png
+sleeping_2/frame_0040_img.png
+sleeping_2/frame_0041_img.png
+sleeping_2/frame_0042_img.png
+sleeping_2/frame_0043_img.png
+sleeping_2/frame_0044_img.png
+sleeping_2/frame_0045_img.png
+sleeping_2/frame_0046_img.png
+sleeping_2/frame_0047_img.png
+sleeping_2/frame_0048_img.png
+sleeping_2/frame_0049_img.png
+sleeping_2/frame_0050_img.png
+temple_2/frame_0001_img.png
+temple_2/frame_0002_img.png
+temple_2/frame_0003_img.png
+temple_2/frame_0004_img.png
+temple_2/frame_0005_img.png
+temple_2/frame_0006_img.png
+temple_2/frame_0007_img.png
+temple_2/frame_0008_img.png
+temple_2/frame_0009_img.png
+temple_2/frame_0010_img.png
+temple_2/frame_0011_img.png
+temple_2/frame_0012_img.png
+temple_2/frame_0013_img.png
+temple_2/frame_0014_img.png
+temple_2/frame_0015_img.png
+temple_2/frame_0016_img.png
+temple_2/frame_0017_img.png
+temple_2/frame_0018_img.png
+temple_2/frame_0019_img.png
+temple_2/frame_0020_img.png
+temple_2/frame_0021_img.png
+temple_2/frame_0022_img.png
+temple_2/frame_0023_img.png
+temple_2/frame_0024_img.png
+temple_2/frame_0025_img.png
+temple_2/frame_0026_img.png
+temple_2/frame_0027_img.png
+temple_2/frame_0028_img.png
+temple_2/frame_0029_img.png
+temple_2/frame_0030_img.png
+temple_2/frame_0031_img.png
+temple_2/frame_0032_img.png
+temple_2/frame_0033_img.png
+temple_2/frame_0034_img.png
+temple_2/frame_0035_img.png
+temple_2/frame_0036_img.png
+temple_2/frame_0037_img.png
+temple_2/frame_0038_img.png
+temple_2/frame_0039_img.png
+temple_2/frame_0040_img.png
+temple_2/frame_0041_img.png
+temple_2/frame_0042_img.png
+temple_2/frame_0043_img.png
+temple_2/frame_0044_img.png
+temple_2/frame_0045_img.png
+temple_2/frame_0046_img.png
+temple_2/frame_0047_img.png
+temple_2/frame_0048_img.png
+temple_2/frame_0049_img.png
+temple_2/frame_0050_img.png
+temple_3/frame_0001_img.png
+temple_3/frame_0002_img.png
+temple_3/frame_0003_img.png
+temple_3/frame_0004_img.png
+temple_3/frame_0005_img.png
+temple_3/frame_0006_img.png
+temple_3/frame_0007_img.png
+temple_3/frame_0008_img.png
+temple_3/frame_0009_img.png
+temple_3/frame_0010_img.png
+temple_3/frame_0011_img.png
+temple_3/frame_0012_img.png
+temple_3/frame_0013_img.png
+temple_3/frame_0014_img.png
+temple_3/frame_0015_img.png
+temple_3/frame_0016_img.png
+temple_3/frame_0017_img.png
+temple_3/frame_0018_img.png
+temple_3/frame_0019_img.png
+temple_3/frame_0020_img.png
+temple_3/frame_0021_img.png
+temple_3/frame_0022_img.png
+temple_3/frame_0023_img.png
+temple_3/frame_0024_img.png
+temple_3/frame_0025_img.png
+temple_3/frame_0026_img.png
+temple_3/frame_0027_img.png
+temple_3/frame_0028_img.png
+temple_3/frame_0029_img.png
+temple_3/frame_0030_img.png
+temple_3/frame_0031_img.png
+temple_3/frame_0032_img.png
+temple_3/frame_0033_img.png
+temple_3/frame_0034_img.png
+temple_3/frame_0035_img.png
+temple_3/frame_0036_img.png
+temple_3/frame_0037_img.png
+temple_3/frame_0038_img.png
+temple_3/frame_0039_img.png
+temple_3/frame_0040_img.png
+temple_3/frame_0041_img.png
+temple_3/frame_0042_img.png
+temple_3/frame_0043_img.png
+temple_3/frame_0044_img.png
+temple_3/frame_0045_img.png
+temple_3/frame_0046_img.png
+temple_3/frame_0047_img.png
+temple_3/frame_0048_img.png
+temple_3/frame_0049_img.png
+temple_3/frame_0050_img.png

FE2E/infer/image_utils.py

@@ -0,0 +1,158 @@
+from PIL import Image
+import matplotlib
+import numpy as np
+
+from typing import List, Union
+import PIL.Image
+
+import torch
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms.functional import resize
+
+def concatenate_images(*image_lists):
+    # Ensure at least one image list is provided
+    if not image_lists or not image_lists[0]:
+        raise ValueError("At least one non-empty image list must be provided")
+    
+    # Determine the maximum width of any single row and the total height
+    max_width = 0
+    total_height = 0
+    row_widths = []
+    row_heights = []
+
+    # Compute dimensions for each row
+    for image_list in image_lists:
+        if image_list:  # Ensure the list is not empty
+            width = sum(img.width for img in image_list)
+            height = image_list[0].height  # Assuming all images in the list have the same height
+            max_width = max(max_width, width)
+            total_height += height
+            row_widths.append(width)
+            row_heights.append(height)
+    
+    # Create a new image to concatenate everything into
+    new_image = Image.new('RGB', (max_width, total_height))
+    
+    # Concatenate each row of images
+    y_offset = 0
+    for i, image_list in enumerate(image_lists):
+        x_offset = 0
+        for img in image_list:
+            new_image.paste(img, (x_offset, y_offset))
+            x_offset += img.width
+        y_offset += row_heights[i]  # Move the offset down to the next row
+    
+    return new_image
+
+
+def colorize_depth_map(depth, mask=None, reverse_color=False):
+    cm = matplotlib.colormaps["Spectral"]
+    # normalize
+    depth = ((depth - depth.min()) / (depth.max() - depth.min()))
+    # colorize
+    if reverse_color:
+        img_colored_np = cm(1 - depth, bytes=False)[:, :, 0:3]  # Invert the depth values before applying colormap
+    else:
+        img_colored_np = cm(depth, bytes=False)[:, :, 0:3] # (h,w,3)
+
+    depth_colored = (img_colored_np * 255).astype(np.uint8) 
+    if mask is not None:
+        masked_image = np.zeros_like(depth_colored)
+        masked_image[mask.numpy()] = depth_colored[mask.numpy()]
+        depth_colored_img = Image.fromarray(masked_image)
+    else:
+        depth_colored_img = Image.fromarray(depth_colored)
+    return depth_colored_img
+
+
+def resize_max_res(
+    img: torch.Tensor,
+    max_edge_resolution: int,
+    resample_method: InterpolationMode = InterpolationMode.BILINEAR,
+) -> torch.Tensor:
+    """
+    Resize image to limit maximum edge length while keeping aspect ratio.
+
+    Args:
+        img (`torch.Tensor`):
+            Image tensor to be resized. Expected shape: [B, C, H, W]
+        max_edge_resolution (`int`):
+            Maximum edge length (pixel).
+        resample_method (`PIL.Image.Resampling`):
+            Resampling method used to resize images.
+
+    Returns:
+        `torch.Tensor`: Resized image.
+    """
+    assert 4 == img.dim(), f"Invalid input shape {img.shape}"
+
+    original_height, original_width = img.shape[-2:]
+    downscale_factor = min(
+        max_edge_resolution / original_width, max_edge_resolution / original_height
+    )
+
+    new_width = int(original_width * downscale_factor)
+    new_height = int(original_height * downscale_factor)
+
+    resized_img = resize(img, (new_height, new_width), resample_method, antialias=True)
+    return resized_img
+
+def resize_back(
+        img: Union[torch.Tensor, np.ndarray, PIL.Image.Image, List[PIL.Image.Image]],
+        target_size: Union[int, tuple[int, int]],
+        resample_method: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
+) -> Union[torch.Tensor, np.ndarray, PIL.Image.Image, List[PIL.Image.Image]]:
+    """
+    Resize image to target size.
+
+    Args:
+        img (`Union[torch.Tensor, np.ndarray, PIL.Image.Image, List[PIL.Image.Image]]`):
+            Image to be resized. 
+        target_size (`Union[int, tuple[int, int]]`):
+            Target size of the resized image.
+        resample_method (`Union[InterpolationMode, int]`):
+            Resampling method used to resize images. 
+
+    Returns:
+        `Union[torch.Tensor, np.ndarray, PIL.Image.Image, List[PIL.Image.Image]]`: Resized image.
+    """
+    if isinstance(img, torch.Tensor): # [B, C, H, W]
+        resized_img = resize(img, target_size, resample_method, antialias=True)
+    if isinstance(img, np.ndarray): # [B, H, W, C]
+        # Conver to Torch.Tensor
+        img = torch.tensor(img).permute(0, 3, 1, 2)
+        resized_img = resize(img, target_size, resample_method, antialias=True)
+        # Convert back to np.ndarray
+        resized_img = resized_img.permute(0, 2, 3, 1).numpy()
+    elif isinstance(img, PIL.Image.Image):
+        target_size = (target_size[1], target_size[0])  # PIL uses (width, height)
+        resized_img = img.resize(target_size, resample_method)
+    elif isinstance(img, list) and all(isinstance(i, PIL.Image.Image) for i in img):
+        target_size = (target_size[1], target_size[0])  # PIL uses (width, height)
+        resized_img = [i.resize(target_size, resample_method) for i in img]
+    return resized_img
+
+def get_pil_resample_method(method_str: str) -> int:
+    resample_method_dict = {
+        "bilinear": Image.BILINEAR,
+        "bicubic": Image.BICUBIC,
+        "nearest": Image.NEAREST,
+    }
+    resample_method = resample_method_dict.get(method_str, None)
+    if resample_method is None:
+        raise ValueError(f"Unknown resampling method: {resample_method}")
+    else:
+        return resample_method
+
+def get_tv_resample_method(method_str: str) -> InterpolationMode:
+    resample_method_dict = {
+        "bilinear": InterpolationMode.BILINEAR,
+        "bicubic": InterpolationMode.BICUBIC,
+        "nearest": InterpolationMode.NEAREST_EXACT,
+        "nearest-exact": InterpolationMode.NEAREST_EXACT,
+    }
+    resample_method = resample_method_dict.get(method_str, None)
+    if resample_method is None:
+        raise ValueError(f"Unknown resampling method: {resample_method}")
+    else:
+        return resample_method

FE2E/infer/inference.py

@@ -0,0 +1,569 @@
+import argparse
+import json
+import itertools
+import math
+import os
+import sys
+import time
+from pathlib import Path
+
+# 添加父目录到系统路径
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import numpy as np
+import torch
+from einops import rearrange, repeat
+from PIL import Image
+from safetensors.torch import load_file
+from torchvision.transforms import functional as F
+from tqdm import tqdm
+import torch.nn.functional as Func
+
+import infer.sampling as sampling
+from modules.autoencoder import AutoEncoder
+from modules.model_edit import Step1XParams, Step1XEdit
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_QWEN_DIR = REPO_ROOT / "Qwen"
+EMPTY_PROMPT_LATENT_PATH = REPO_ROOT / "latent" / "no_info.npz"
+
+
+def cudagc():
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+
+
+def load_state_dict(model, ckpt_path, device="cuda", strict=False, assign=True):
+    if Path(ckpt_path).suffix == ".safetensors":
+        state_dict = load_file(ckpt_path, device)
+    else:
+        state_dict = torch.load(ckpt_path, map_location="cpu")
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=strict, assign=assign)
+    if len(missing) > 0 and len(unexpected) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+        print("\n" + "-" * 79 + "\n")
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+    elif len(missing) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+    elif len(unexpected) > 0:
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+    return model
+
+
+def load_models(dit_path=None, ae_path=None, qwen2vl_model_path=None, device="cuda", max_length=256, dtype=torch.bfloat16, args=None):
+    empty_llm = args is not None and hasattr(args, 'prompt_type') and args.prompt_type == 'empty'
+    if empty_llm:
+        print("[INFO] prompt_type=empty, 跳过Qwen模型加载")
+        qwen2vl_encoder = None
+    else:
+        # Lazy import to avoid pulling transformers/vision stack during evaluation with prompt_type=empty.
+        from modules.conditioner import Qwen25VL_7b_Embedder as Qwen2VLEmbedder
+        qwen2vl_encoder = Qwen2VLEmbedder(
+            qwen2vl_model_path,
+            device=device,
+            max_length=max_length,
+            dtype=dtype,
+            args=args,
+        )
+
+    with torch.device("meta"):
+        ae = AutoEncoder(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        )
+
+        step1x_params = Step1XParams(
+            in_channels=64,
+            out_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+        )
+        dit = Step1XEdit(step1x_params)
+
+    ae = load_state_dict(ae, ae_path, 'cpu')
+    dit = load_state_dict(dit, dit_path, 'cpu')
+
+    ae = ae.to(dtype=torch.float32)
+
+    return ae, dit, qwen2vl_encoder
+
+
+def equip_dit_with_lora_sd_scripts(ae, text_encoders, dit, lora, device='cuda'):
+    from safetensors.torch import load_file
+    weights_sd = load_file(lora)
+    is_lora = True
+    from library import lora_module
+    module = lora_module
+    lora_model, _ = module.create_network_from_weights(1.0, None, ae, text_encoders, dit, weights_sd, True)
+    lora_model.merge_to(text_encoders, dit, weights_sd)
+
+    lora_model.set_multiplier(1.0)
+    return lora_model
+
+class ImageGenerator:
+
+    def __init__(
+        self,
+        dit_path=None,
+        ae_path=None,
+        qwen2vl_model_path=None,
+        device="cuda",
+        max_length=640,
+        dtype=torch.bfloat16,
+        quantized=False,
+        offload=False,
+        lora=None,
+        args=None,
+    ) -> None:
+        self.device = torch.device(device)
+        self.args = args
+        self.ae, self.dit, self.llm_encoder = load_models(
+            dit_path=dit_path,
+            ae_path=ae_path,
+            qwen2vl_model_path=qwen2vl_model_path,
+            max_length=max_length,
+            dtype=dtype,
+            device=device,
+            args=args,
+        )
+        if not quantized:
+            self.dit = self.dit.to(dtype=torch.bfloat16)
+        else:
+            self.dit = self.dit.to(dtype=torch.float8_e4m3fn)
+        if not offload:
+            self.dit = self.dit.to(device=self.device)
+            self.ae = self.ae.to(device=self.device)
+        self.quantized = quantized
+        self.offload = offload
+        if lora is not None:
+            self.lora_module = equip_dit_with_lora_sd_scripts(
+                self.ae,
+                [self.llm_encoder],
+                self.dit,
+                lora,
+                device=self.dit.device,
+            )
+        else:
+            self.lora_module = None
+
+    def prepare(self, prompt, img, ref_image, ref_image_raw, empty_llm=False):
+        bs, _, h, w = img.shape
+        bs, _, ref_h, ref_w = ref_image.shape
+
+        assert h == ref_h and w == ref_w
+
+        if bs == 1 and not isinstance(prompt, str):
+            bs = len(prompt)
+        elif bs >= 1 and isinstance(prompt, str):
+            prompt = [prompt] * bs
+
+        img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)  #2,16,82,110->2,2255,64
+        ref_img = rearrange(ref_image, "b c (ref_h ph) (ref_w pw) -> b (ref_h ref_w) (c ph pw)", ph=2, pw=2)  # 将二维图像"压平"成一维序列 这是为 Transformer 模型准备的，因为它处理的是序列数据
+        if img.shape[0] == 1 and bs > 1:
+            img = repeat(img, "1 ... -> bs ...", bs=bs)
+            ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
+        #img 和 ref_img 已经不再是二维的图片了，而是变成了一个 "patches" (图像块) 的序列。一个块是64维度的。Transformer不知道这2255个图像块哪个在左上角，哪个在右下角。
+        img_ids = torch.zeros(h // 2, w // 2, 3)  #41,55,3 # h 和 w 是潜在空间的高和宽，但 rearrange 操作把 2x2 的小块合并了# 所以实际的网格大小是 h/2 x w/2# 最后的 3 代表每个坐标有3个分量 (一个预留, Y坐标, X坐标)
+        img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
+        img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]  #通过广播机制，第 i 行的所有点的第二个分量都被赋值为 i
+        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)  #将二维坐标网格"压平"成一维序列，并复制到对应的batch size
+
+        ref_img_ids = torch.zeros(ref_h // 2, ref_w // 2, 3)
+
+        ref_img_ids[..., 1] = ref_img_ids[..., 1] + torch.arange(ref_h // 2)[:, None]
+        ref_img_ids[..., 2] = ref_img_ids[..., 2] + torch.arange(ref_w // 2)[None, :]
+        ref_img_ids = repeat(ref_img_ids, "ref_h ref_w c -> b (ref_h ref_w) c", b=bs)
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        if self.offload:
+            self.llm_encoder = self.llm_encoder.to(self.device)
+
+        if empty_llm:
+            empty_prompt_cache = getattr(self.args, "empty_prompt_cache", None) if self.args is not None else None
+            cache_path = Path(empty_prompt_cache) if empty_prompt_cache else EMPTY_PROMPT_LATENT_PATH
+            data = np.load(cache_path)
+            txt = torch.from_numpy(data['embeds']).to(img.device).unsqueeze(0)
+            txt = torch.cat([txt, txt], dim=0)
+            mask = torch.from_numpy(data['masks']).to(img.device).unsqueeze(0)
+            mask = torch.cat([mask, mask], dim=0)
+        else:
+            txt, mask = self.llm_encoder(prompt, ref_image_raw)  #之所以都要复制一份，是因为有正负两种prompt
+
+        if self.offload:
+            self.llm_encoder = self.llm_encoder.cpu()
+            cudagc()
+
+        txt_ids = torch.zeros(bs, txt.shape[1], 3)
+
+        img = torch.cat([img, ref_img.to(device=img.device, dtype=img.dtype)], dim=-2)  #2,4550,64 在patch上concat？？？
+        img_ids = torch.cat([img_ids, ref_img_ids], dim=-2)
+
+        return {
+            "img": img,
+            "mask": mask,
+            "img_ids": img_ids.to(img.device),  #图像坐标
+            "llm_embedding": txt.to(img.device),  #文字向量
+            "txt_ids": txt_ids.to(img.device),  #文字坐标
+        }
+
+    @staticmethod
+    def process_diff_norm(diff_norm, k):
+        pow_result = torch.pow(diff_norm, k)
+
+        result = torch.where(
+            diff_norm > 1.0,
+            pow_result,
+            torch.where(diff_norm < 1.0, torch.ones_like(diff_norm), diff_norm),
+        )
+        return result
+
+    def denoise(
+        self,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+        llm_embedding: torch.Tensor,
+        txt_ids: torch.Tensor,
+        timesteps: list[float],
+        cfg_guidance: float = 6.0,
+        mask=None,
+        show_progress=False,
+        timesteps_truncate=1.0,
+    ):
+        if self.offload:
+            self.dit = self.dit.to(self.device)
+        if show_progress:
+            pbar = tqdm(itertools.pairwise(timesteps), desc='denoising...')
+        else:
+            pbar = itertools.pairwise(timesteps)
+        '''
+        Cond 0 RGB
+        Uncd 0 RGB
+        '''
+        for t_curr, t_prev in pbar:
+            '''
+            若输入维度是2，无所谓，维度是1则：
+            imgN D RGB
+            imgN D RGB
+            '''
+            if img.shape[0] == 1 and cfg_guidance != -1:
+                img = torch.cat([img, img], dim=0)
+            t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+
+            pred, feat = self.dit(
+                img=img,
+                img_ids=img_ids,
+                txt_ids=txt_ids,
+                timesteps=t_vec,
+                llm_embedding=llm_embedding,
+                t_vec=t_vec,
+                mask=mask,
+            )
+
+            assert cfg_guidance != -1, " cfg_guidance must not be -1 NOW!!!!"
+            cond, uncond = (
+                pred[0:pred.shape[0] // 2, :],
+                pred[pred.shape[0] // 2:, :],
+            )
+            '''
+            Cond D ??? <- pred
+            Uncd D ???
+            '''
+            pred = uncond + cfg_guidance * (cond - uncond)  #1,4608,64
+            pred1 = cond  #todo only support single denoise!!!
+            '''
+                Cond 0 RGB
+              + pred D ???
+                temI D ??? 
+            '''
+            tem_img = img[0:img.shape[0] // 2, :] + (t_prev - t_curr) * pred  #1,4608,64
+            img_input_length = img.shape[1] // 2
+            '''
+                tmpI [D](√)  ???(x)
+            cat Cond  0(x)   [RGB](√)
+                imgN  [D]    [RGB]
+            '''
+            img = torch.cat(
+                [
+                    tem_img[:, :img_input_length],  #1,2304,64
+                    img[:img.shape[0] // 2, img_input_length:],  #1,2304,64
+                ],
+                dim=1)  #1,4608,64
+        if self.offload:
+            self.dit = self.dit.cpu()
+            cudagc()
+
+        return img[:, :img.shape[1] // 2], pred1[:, img.shape[1] // 2:]
+
+    def double_denoise(self,img,img_ids,llm_embedding,txt_ids,timesteps,cfg_guidance=6.0,mask=None,height=None,width=None):
+        if img.shape[0] == 1 and cfg_guidance != -1:
+            img = torch.cat([img, img], dim=0)
+        
+        t_vec = torch.full((img.shape[0],), 1.0, dtype=img.dtype, device=img.device)
+
+        pred, _ = self.dit(
+            img=img,
+            img_ids=img_ids,
+            txt_ids=txt_ids,
+            timesteps=t_vec,
+            llm_embedding=llm_embedding,
+            t_vec=t_vec,
+            mask=mask,
+        )
+
+        assert cfg_guidance != -1, " cfg_guidance must not be -1 NOW!!!!"
+        pred, uncond = (
+            pred[0:pred.shape[0] // 2, :],
+            pred[pred.shape[0] // 2:, :],
+        )
+        Lpred,Rpred = self.unpack_latents(pred, height//16, width//16)
+        return Lpred.to(torch.float32),Rpred.to(torch.float32)
+
+    @staticmethod
+    def unpack(x: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+            h=math.ceil(height / 16),
+            w=math.ceil(width / 16),
+            ph=2,
+            pw=2,
+        )
+        
+    @staticmethod
+    def unpack_latents(x: torch.Tensor, packed_latent_height: int, packed_latent_width: int):
+            """
+            x: [b (h w) (c ph pw)] -> [b c (h ph) (w pw)], ph=2, pw=2
+            """
+            import einops
+            x = einops.rearrange(x, "b (p h w) (c ph pw) -> b p c (h ph) (w pw)", h=packed_latent_height, w=packed_latent_width, ph=2, pw=2, p=2)
+            return x[:, 0], x[:, 1]
+        
+    @staticmethod
+    def load_image(image):
+        from PIL import Image
+
+        if isinstance(image, np.ndarray):
+            image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
+            image = image.unsqueeze(0)
+            return image
+        elif isinstance(image, Image.Image):
+            image = F.to_tensor(image.convert("RGB"))
+            image = image.unsqueeze(0)
+            return image
+        elif isinstance(image, torch.Tensor):
+            return image
+        elif isinstance(image, str):
+            image = F.to_tensor(Image.open(image).convert("RGB"))
+            image = image.unsqueeze(0)
+            return image
+        else:
+            raise ValueError(f"Unsupported image type: {type(image)}")
+
+    def output_process_image(self, resize_img, image_size):
+        res_image = resize_img.resize(image_size)
+        return res_image
+
+    def input_process_image(self, img):
+        if isinstance(img, torch.Tensor):
+            w, h = img.shape[-1], img.shape[-2]
+        elif isinstance(img, Image.Image):
+            w, h = img.size
+
+        if w <= 1024 and h <= 768:
+            w_new, h_new = 1024, 768
+        elif w <= 1280 and h <= 960:
+            w_new, h_new = 1216, 352
+        elif w <= 6048 and h <= 4032:
+            w_new, h_new = 864, 576
+        else:
+            w_new, h_new = w, h
+
+        if isinstance(img, torch.Tensor):
+            img_resized = Func.interpolate(img, (h_new, w_new), mode='bilinear', align_corners=False)
+            img_resized = img_resized.clamp(0, 1)
+        else:
+            img_resized = img.resize((w_new, h_new))
+
+        return img_resized, (w_new, h_new)
+
+    @torch.inference_mode()
+    def generate_image(
+        self,prompt,negative_prompt,ref_images,num_steps,cfg_guidance,seed,num_samples=1,init_image=None,image2image_strength=0.0,show_progress=False,size_level=512,args=None,judge=None,name=None
+    ):
+        assert num_samples == 1, "num_samples > 1 is not supported yet."
+
+        ref_images_raw, img_info = self.input_process_image(ref_images)
+        if isinstance(ref_images, Image.Image):
+            ref_images_raw = self.load_image(ref_images_raw)
+
+        height, width = ref_images_raw.shape[-2], ref_images_raw.shape[-1]
+
+        ref_images_raw = ref_images_raw.to(self.device)
+        if self.offload:
+            self.ae = self.ae.to(self.device)
+        ref_images = self.ae.encode(ref_images_raw.to(self.device) * 2 - 1)  #bs,3,656,880 -> 1,16,82,110
+        #加入cache
+        
+        if self.offload:
+            self.ae = self.ae.cpu()
+            cudagc()
+
+        seed = int(seed)
+        seed = torch.Generator(device="cpu").seed() if seed < 0 else seed
+
+        t0 = time.perf_counter()
+
+        if init_image is not None:
+            init_image = self.load_image(init_image)
+            init_image = init_image.to(self.device)
+            init_image = torch.nn.functional.interpolate(init_image, (height, width))
+            if self.offload:
+                self.ae = self.ae.to(self.device)
+            init_image = self.ae.encode(init_image.to() * 2 - 1)
+            if self.offload:
+                self.ae = self.ae.cpu()
+                cudagc()
+
+        if args is not None and hasattr(args, 'single_denoise') and not args.single_denoise:
+            x = torch.randn(num_samples,16,height // 8,width // 8,device=self.device,dtype=torch.bfloat16,generator=torch.Generator(device=self.device).manual_seed(seed),)
+        else:
+            x= torch.zeros(num_samples,16,height // 8,width // 8,device=self.device,dtype=torch.bfloat16,)
+
+
+        timesteps = sampling.get_schedule(num_steps, x.shape[-1] * x.shape[-2] // 4, shift=True)
+
+        if init_image is not None:
+            t_idx = int((1 - image2image_strength) * num_steps)
+            t = timesteps[t_idx]
+            timesteps = timesteps[t_idx:]
+            x = t * x + (1.0 - t) * init_image.to(x.dtype)
+
+        x = torch.cat([x, x], dim=0)
+        ref_images = torch.cat([ref_images, ref_images], dim=0)  #这里是为了有无prompt
+        ref_images_raw = torch.cat([ref_images_raw, ref_images_raw], dim=0)
+        
+        # 检查args和prompt_type属性
+        empty_llm = args is not None and hasattr(args, 'prompt_type') and args.prompt_type == 'empty'
+        
+        inputs = self.prepare(
+            [prompt, negative_prompt],
+            x,  #img这个gt给的是全噪声在推理
+            ref_image=ref_images,
+            ref_image_raw=ref_images_raw,
+            empty_llm=empty_llm)
+
+        with torch.autocast(device_type=self.device.type,dtype=torch.bfloat16):
+            # Lpred,Rpred = self.double_denoise(**inputs,cfg_guidance=cfg_guidance,timesteps=timesteps,height=height,width=width)#图像中包括ref image
+            Lpred,Rpred = self.denoise(**inputs,cfg_guidance=cfg_guidance,timesteps=timesteps,show_progress=show_progress,timesteps_truncate=1.0,)#图像中包括ref image
+            Lpred=self.unpack(Lpred.float(),height,width)
+            Rpred=self.unpack(Rpred.float(),height,width)
+        if judge is not None:
+            judge = Func.interpolate(judge, (height, width), mode='bilinear', align_corners=False)
+            training_gt=self.ae.encode(judge)
+            traing_loss = torch.nn.functional.mse_loss(Rpred,training_gt)
+            print(f"training_loss with rgb2: {traing_loss}")
+            
+            norm = torch.linalg.norm(judge, dim=1, keepdim=True)
+            norm[norm < 1e-9] = 1e-9
+            judge = judge / norm
+            training_gt =self.ae.encode(judge)
+            training_loss = torch.nn.functional.mse_loss(Rpred,training_gt)
+            print(f"training_loss with normed_rgb: {training_loss}")
+        Lpred = self.ae.decode(Lpred)
+        Rpred = self.ae.decode(Rpred)
+        
+            
+        Lpred = Lpred.clamp(-1, 1)
+        Lpred = Lpred.mul(0.5).add(0.5)
+        Rpred = Rpred.clamp(-1, 1)
+        # Rpred = Rpred.mul(0.5).add(0.5)
+
+        images_list = []
+        for img in Rpred.float():
+            images_list.append(self.output_process_image(F.to_pil_image(img), img_info))
+        return images_list, Lpred.float(), Rpred.float()
+
+
+def main():
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_path', type=str, required=True, help='Path to the model checkpoint')
+    parser.add_argument('--input_dir', type=str, required=True, help='Path to the input image directory')
+    parser.add_argument('--output_dir', type=str, required=True, help='Path to the output image directory')
+    parser.add_argument('--json_path', type=str, required=True, help='Path to the JSON file containing image names and prompts')
+    parser.add_argument('--seed', type=int, default=42, help='Random seed for generation')
+    parser.add_argument('--num_steps', type=int, default=28, help='Number of diffusion steps')
+    parser.add_argument('--cfg_guidance', type=float, default=6.0, help='CFG guidance strength')
+    parser.add_argument('--size_level', default=512, type=int)
+    parser.add_argument('--offload', action='store_true', help='Use offload for large models')
+    parser.add_argument('--quantized', action='store_true', help='Use fp8 model weights')
+    parser.add_argument('--lora', type=str, default=None)
+    parser.add_argument('--qwen2vl_model_path', type=str, default=str(DEFAULT_QWEN_DIR), help='Path to the local Qwen2.5-VL model directory')
+    parser.add_argument('--empty_prompt_cache', type=str, default=str(EMPTY_PROMPT_LATENT_PATH), help='Path to the empty-prompt latent cache')
+    args = parser.parse_args()
+
+    assert os.path.exists(args.input_dir), f"Input directory {args.input_dir} does not exist."
+    assert os.path.exists(args.json_path), f"JSON file {args.json_path} does not exist."
+
+    args.output_dir = args.output_dir.rstrip('/') + ('-offload' if args.offload else "") + ('-quantized' if args.quantized else "") + f"-{args.size_level}"
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    image_and_prompts = json.load(open(args.json_path, 'r'))
+
+    image_edit = ImageGenerator(
+        ae_path=os.path.join(args.model_path, 'vae.safetensors'),
+        dit_path=os.path.join(args.model_path, "step1x-edit-i1258-FP8.safetensors" if args.quantized else "step1x-edit-i1258.safetensors"),
+        qwen2vl_model_path=args.qwen2vl_model_path,
+        max_length=640,
+        quantized=args.quantized,
+        offload=args.offload,
+        lora=args.lora,
+    )
+
+    time_list = []
+    for image_name, prompt in image_and_prompts.items():
+        image_path = os.path.join(args.input_dir, image_name)
+        output_path = os.path.join(args.output_dir, image_name)
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        start_time = time.time()
+
+        images, _, _ = image_edit.generate_image(
+            prompt,
+            negative_prompt="",
+            ref_images=Image.open(image_path).convert("RGB"),
+            num_samples=1,
+            num_steps=args.num_steps,
+            cfg_guidance=args.cfg_guidance,
+            seed=args.seed,
+            show_progress=True,
+            size_level=args.size_level,
+        )
+
+        print(f"Time taken: {time.time() - start_time:.2f} seconds")
+        time_list.append(time.time() - start_time)
+
+        images[0].save(output_path, lossless=True)
+    if len(time_list) > 1:
+        print(f'average time for {args.output_dir}: ', sum(time_list[1:]) / len(time_list[1:]))
+
+
+if __name__ == "__main__":
+    main()

FE2E/infer/inner_evaluation.py

@@ -0,0 +1,598 @@
+import logging
+import os
+import sys
+import csv  # 保留csv库用于保存结果
+import multiprocessing as mp
+import time
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from tabulate import tabulate
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+import cv2
+from infer.dataset import (
+    BaseDepthDataset,
+    DatasetMode,
+    get_dataset,
+    get_pred_name,
+)
+from .util import metric, normal_utils
+from .util.alignment import (align_depth_least_square, depth2disparity, disparity2depth, depth2log_space, log_space2depth)
+from .util.metric import MetricTracker
+from infer.image_utils import colorize_depth_map
+
+eval_metrics = [
+    "abs_relative_difference",
+    "squared_relative_difference",
+    "rmse_linear",
+    "rmse_log",
+    "delta1_acc",
+    "delta2_acc",
+    "delta3_acc",
+]
+
+
+def save_visualization_worker(save_vis_path, safe_pred_name, cfg_suffix, depth_pred_np, depth_raw_np, valid_mask_np, input_rgb_data, rank):
+    """
+    可视化保存的工作函数，在独立进程中运行
+    Args:
+        save_vis_path: 保存路径
+        safe_pred_name: 安全的预测文件名
+        cfg_suffix: cfg后缀
+        depth_pred_np: 预测深度图 numpy数组
+        depth_raw_np: GT深度图 numpy数组
+        valid_mask_np: 有效掩码 numpy数组
+        input_rgb_data: 输入RGB图像数据
+        rank: GPU rank
+    """
+    try:
+        # 转换为torch tensor用于colorize_depth_map
+        depth_pred_ts = torch.from_numpy(depth_pred_np)
+        depth_raw_ts = torch.from_numpy(depth_raw_np)
+        valid_mask_ts = torch.from_numpy(valid_mask_np)
+        
+        # 1. 保存预测深度图
+        depth_pred_vis = colorize_depth_map(depth_pred_ts)
+        pred_save_path = os.path.join(save_vis_path, f"{safe_pred_name}{cfg_suffix}_pred.png")
+        depth_pred_vis.save(pred_save_path)
+        print(f"saved: {pred_save_path}")
+        # 3. 保存误差图
+        # 计算绝对相对误差
+        abs_rel_error = torch.abs(depth_pred_ts - depth_raw_ts) / (depth_raw_ts + 1e-6)
+        abs_rel_error = abs_rel_error * valid_mask_ts.float()
+        
+        # 使用matplotlib生成误差图
+        import matplotlib
+        matplotlib.use('Agg')  # 使用非交互式后端
+        import matplotlib.pyplot as plt
+        import matplotlib.cm as cm
+        
+        error_np = abs_rel_error.numpy()
+        # 设置误差显示范围
+        vmax = 0.2  # 可以根据需要调整
+        error_normalized = np.clip(error_np / vmax, 0, 1)
+        
+        # 应用颜色映射
+        jet_cmap = cm.get_cmap('jet')
+        error_colored = jet_cmap(error_normalized)[:, :, :3]  # 去掉alpha通道
+        error_colored = (error_colored * 255).astype(np.uint8)
+        
+        # 将无效区域设为黑色
+        error_colored[~valid_mask_np] = [0, 0, 0]
+        
+        error_save_path = os.path.join(save_vis_path, f"{safe_pred_name}{cfg_suffix}_error.png")
+        plt.imsave(error_save_path, error_colored)
+        print(f"saved: {error_save_path}")
+        # 关闭matplotlib figure以释放内存
+        plt.close('all')
+        
+    except Exception as e:
+        print(f"[VIS-Worker-{rank}] 可视化保存失败: {e}", file=sys.stderr)
+
+
+def prepare_input_rgb_data(input_rgb):
+    """
+    预处理输入RGB数据，转换为numpy格式供子进程使用
+    """
+    if input_rgb is None:
+        return None
+        
+    try:
+        # 处理不同格式的输入图像
+        if isinstance(input_rgb, torch.Tensor):
+            # 如果是torch tensor，转换为numpy
+            if input_rgb.dim() == 4:  # Batch dimension
+                input_rgb = input_rgb[0]
+            if input_rgb.dim() == 3 and input_rgb.shape[0] == 3:  # CHW格式
+                input_rgb = input_rgb.permute(1, 2, 0)
+            
+            # 确保值在[0,1]范围内
+            if input_rgb.max() <= 1.0:
+                input_rgb = (input_rgb * 255).clamp(0, 255).byte()
+            
+            input_rgb_np = input_rgb.cpu().numpy()
+            
+        elif isinstance(input_rgb, np.ndarray):
+            input_rgb_np = input_rgb
+            # 确保值在正确范围内
+            if input_rgb_np.max() <= 1.0:
+                input_rgb_np = (input_rgb_np * 255).astype(np.uint8)
+        else:
+            # 假设是PIL图像或其他格式
+            input_rgb_np = np.array(input_rgb)
+            
+        return input_rgb_np.copy()  # 创建副本避免进程间共享问题
+        
+    except Exception as e:
+        print(f"处理输入RGB数据失败: {e}", file=sys.stderr)
+        return None
+
+
+def evaluate_single_prediction(pred_depth, depth_raw, valid_mask, dataset, device, metric_funcs, alignment_max_res=None, save_pred_vis=False, save_vis_path=None, pred_name=None, cfg_suffix="", alignment="least_square", rank=0, input_rgb=None):
+    """Args:  pred_depth: 预测的深度图 (numpy array, [0,1])        depth_raw: 真实深度图 (numpy array)         valid_mask: ��效掩码 (numpy array)        dataset: 数据集对象        device: 计算设备        metric_funcs: 评估指标函数列表        alignment_max_res: 对齐时的最大分辨率        save_pred_vis: 是否保存可视化结果        save_vis_path: 可视化保存路径        pred_name: 预测文件名        cfg_suffix: cfg后缀，用于区分不同的cfg设置
+    alignment: 对齐方式，可选值为"least_square"或"least_square_disparity"
+    rank: GPU rank，用于多进程图像保存
+    input_rgb: 输入RGB图像 (torch.Tensor or PIL.Image or numpy.ndarray)
+    Returns: sample_metric: 该样本的所有评估指标列表"""
+    # 确保预测深度图的维度正确
+    if len(pred_depth.shape) == 3:
+        pred_depth = pred_depth.mean(0)  # [0,1]
+
+    # 调整预测深度图尺寸以匹配真实深度图
+    if pred_depth.shape != depth_raw.shape:
+        pred_depth = cv2.resize(pred_depth, (depth_raw.shape[1], depth_raw.shape[0]), interpolation=cv2.INTER_LINEAR)
+
+    if "least_square" == alignment:
+        depth_pred, scale, shift = align_depth_least_square(
+            gt_arr=depth_raw,
+            pred_arr=pred_depth,
+            valid_mask_arr=valid_mask,
+            return_scale_shift=True,
+            max_resolution=alignment_max_res,
+        )
+    elif "log_space" == alignment:
+        gt_log, gt_non_neg_mask = depth2log_space(depth=depth_raw, return_mask=True)
+        pred_non_neg_mask = pred_depth > 0
+        valid_nonnegative_mask = valid_mask & gt_non_neg_mask & pred_non_neg_mask
+
+        # 确保输入是numpy数组类型
+        if isinstance(gt_log, torch.Tensor):
+            gt_log = gt_log.cpu().numpy()
+
+        log_space_pred, scale, shift = align_depth_least_square(
+            gt_arr=gt_log,
+            pred_arr=pred_depth,
+            valid_mask_arr=valid_nonnegative_mask,
+            return_scale_shift=True,
+            max_resolution=alignment_max_res,
+        )
+        log_space_pred = np.clip(log_space_pred, a_min=None, a_max=5.)
+        depth_pred = log_space2depth(log_space_pred)
+    # 裁剪到数据集的深度范围
+    depth_pred = np.clip(depth_pred, a_min=dataset.min_depth, a_max=dataset.max_depth)
+
+    # 裁剪到 d > 0 以便评估
+    depth_pred = np.clip(depth_pred, a_min=1e-6, a_max=None)
+
+    # 转换到设备进行评估
+    depth_pred_ts = torch.from_numpy(depth_pred).to(device)
+    depth_raw_ts = torch.from_numpy(depth_raw).to(device)
+    valid_mask_ts = torch.from_numpy(valid_mask).to(device)
+
+    # 启动可视化保存进程（同步）
+    if save_pred_vis and save_vis_path is not None and pred_name is not None:
+        safe_pred_name = pred_name.replace('/', '_').replace('\\', '_')
+        input_rgb_data = prepare_input_rgb_data(input_rgb)
+        vis_process = mp.Process(
+            target=save_visualization_worker,
+            args=(
+                save_vis_path,
+                safe_pred_name,
+                cfg_suffix,
+                depth_pred.copy(),
+                depth_raw.copy(),
+                valid_mask.copy(),
+                input_rgb_data,
+                rank
+            )
+        )
+        vis_process.start()
+        # save_visualization_worker(save_vis_path, safe_pred_name, cfg_suffix, depth_pred.copy(), depth_raw.copy(), valid_mask.copy(), input_rgb_data, rank)
+
+    # 计算评估指标
+    sample_metric = []
+    for met_func in metric_funcs:
+        _metric = met_func(depth_pred_ts, depth_raw_ts, valid_mask_ts).item()
+        sample_metric.append(_metric)
+
+    return sample_metric
+
+
+def evaluation_depth_custom_parallel(rank, world_size, output_dir, dataset_config, args, pipeline, base_data_dir, pred_suffix="", alignment="least_square", alignment_max_res=None, prediction_dir=None, save_pred_vis=False):
+    """
+    支持多GPU并行的深度评估函数
+    """
+    import time
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    cuda_avail = torch.cuda.is_available()
+    device = torch.device(f"cuda:{rank}")
+
+    cfg_data = OmegaConf.load(dataset_config)
+
+    dataset: BaseDepthDataset = get_dataset(cfg_data, base_data_dir=base_data_dir, mode=DatasetMode.EVAL, prompt_type=args.prompt_type)
+
+    # 获取数据集名称，用于CSV表命名
+    dataset_name = dataset.__class__.__name__
+
+    # 初始化存储结果的数据列表
+    results_data = []
+
+    # 计算每个GPU处理的数据范围
+    total_samples = len(dataset)
+    if args.num_samples > 0:
+        total_samples = min(args.num_samples, total_samples)
+
+    chunk_size = total_samples // world_size
+    start_idx = rank * chunk_size
+    end_idx = start_idx + chunk_size if rank < world_size - 1 else total_samples
+
+    from torch.utils.data import SubsetRandomSampler
+    indices = list(range(start_idx, end_idx))
+
+    dataloader = DataLoader(dataset, batch_size=1, num_workers=0 if args.debug else 4, pin_memory=True, sampler=SubsetRandomSampler(indices),shuffle=False)
+
+    metric_funcs = [getattr(metric, _met) for _met in eval_metrics]
+
+    # 为cfg=1和cfg=6分别创建metric tracker
+    metric_tracker_Lpred = MetricTracker(*[m.__name__ for m in metric_funcs])
+    metric_tracker_Lpred.reset()
+    metric_tracker_Rpred = MetricTracker(*[m.__name__ for m in metric_funcs])
+    metric_tracker_Rpred.reset()
+
+    if save_pred_vis:
+        save_vis_path = os.path.join(output_dir, "vis")
+        os.makedirs(save_vis_path, exist_ok=True)
+
+        # 创建CSV保存目录 - 每个卡都创建
+        csv_save_path = os.path.join(output_dir, "csv_results")
+        os.makedirs(csv_save_path, exist_ok=True)
+    else:
+        save_vis_path = None
+        csv_save_path = None
+
+    processing_times = []
+    vis_processes = []  # 用于管理可视化进程
+    max_vis_processes = 4  # 限制同时运行的可视化进程数量
+
+    sample_count = 0
+    for data in dataloader:
+        sample_count += 1
+
+        depth_raw_ts = data["depth_raw_linear"].squeeze()
+        valid_mask_ts = data["valid_mask_raw"].squeeze()
+        rgb_name = data["rgb_relative_path"][0]
+
+        depth_raw = depth_raw_ts.numpy()
+        valid_mask = valid_mask_ts.numpy()
+
+        # Get predictions
+        rgb_basename = os.path.basename(rgb_name)
+        pred_basename = get_pred_name(rgb_basename, dataset.name_mode, suffix=pred_suffix)
+        pred_name = os.path.join(os.path.dirname(rgb_name), pred_basename)
+
+        start_time = time.time()
+        image_list, Lpred, Rpred = pipeline.generate_image(args.prompt if args.prompt_type == "query" else data["prompt"][0], negative_prompt="", ref_images=data["rgb"], num_samples=1, num_steps=args.num_steps, cfg_guidance=args.cfg_guidance, seed=args.seed + rank, show_progress=False, size_level=args.size_level, args=args)
+        end_time = time.time()
+        processing_times.append(end_time - start_time)
+
+        Lpred = Lpred[0].cpu().numpy()
+        
+        # 保存可视化结果（使用新的多进程方式）
+        if save_pred_vis and save_vis_path is not None:
+            # 清理文件名，替换路径分隔符
+            safe_pred_name = pred_name.replace('/', '_').replace('\\', '_')
+            
+            # 预处理输入RGB数据
+            input_rgb_data = prepare_input_rgb_data(data["rgb"])
+            
+            # 限制同时运行的可视化进程数量
+            while len([p for p in vis_processes if p.is_alive()]) >= max_vis_processes:
+                # 等待一些进程完成
+                for p in vis_processes[:]:
+                    if not p.is_alive():
+                        vis_processes.remove(p)
+                if len([p for p in vis_processes if p.is_alive()]) >= max_vis_processes:
+                    time.sleep(0.1)  # 短暂等待
+            
+            # 创建子进程进行可视化保存
+            vis_process = mp.Process(
+                target=save_visualization_worker,
+                args=(
+                    save_vis_path,
+                    safe_pred_name,
+                    "_Lpred",
+                    Lpred.copy(),
+                    depth_raw.copy(),
+                    valid_mask.copy(),
+                    input_rgb_data,
+                    rank
+                )
+            )
+            vis_process.start()
+            vis_processes.append(vis_process)
+        
+        sample_metric_Lpred = evaluate_single_prediction(pred_depth=Lpred, depth_raw=depth_raw, valid_mask=valid_mask, dataset=dataset, device=device, metric_funcs=metric_funcs, alignment_max_res=alignment_max_res, save_pred_vis=False, save_vis_path=None, pred_name=pred_name, cfg_suffix="_Lpred", alignment=alignment, rank=rank, input_rgb=data["rgb"])
+
+        for i, met_func in enumerate(metric_funcs):
+            metric_name = met_func.__name__
+            metric_tracker_Lpred.update(metric_name, sample_metric_Lpred[i])
+
+        # 输出每个样本的结果
+        img_id = os.path.basename(rgb_name).replace('.png', '').replace('.jpg', '')
+        global_sample_idx = start_idx + sample_count
+
+        # CFG=1结果
+        abs_rel_Lpred = sample_metric_Lpred[0]  # abs_relative_difference
+        rmse_Lpred = sample_metric_Lpred[2]  # rmse_linear
+        delta1_Lpred = sample_metric_Lpred[4]  # delta1_acc
+
+        # 修改输出格式
+        if args.save_viz:
+            print(f"|{global_sample_idx:03d}|{abs_rel_Lpred:.4f}|{rmse_Lpred:.4f}|{delta1_Lpred:.4f}|", file=sys.stderr)
+        elif not args.save_viz:
+            print(f"[GPU:{rank}] 样本:{global_sample_idx:03d}/{total_samples} | ID:{img_id:<12}", file=sys.stderr)
+            print(f"  CFG=1: abs_rel:{abs_rel_Lpred:.4f} | rmse:{rmse_Lpred:.4f} | a1:{delta1_Lpred:.4f}", file=sys.stderr)
+            print(f"  时间: {processing_times[-1]:.2f}s", file=sys.stderr)
+
+        # 所有卡都保存结果到列表
+        if args.save_viz:
+            results_data.append({'GPU_Rank': rank, 'Sample_ID': global_sample_idx, 'Image_Name': rgb_name, 'abs_rel': abs_rel_Lpred, 'rmse': rmse_Lpred, 'delta1': delta1_Lpred, 'processing_time': processing_times[-1]})
+
+    # 等待所有可视化进程完成
+    if save_pred_vis:
+        print(f"[GPU:{rank}] 等待可视化进程完成...", file=sys.stderr)
+        for p in vis_processes:
+            p.join(timeout=30)  # 设置超时时间
+            if p.is_alive():
+                print(f"[GPU:{rank}] 可视化进程超时，强制终止", file=sys.stderr)
+                p.terminate()
+        print(f"[GPU:{rank}] 所有可视化进程已完成", file=sys.stderr)
+
+    if args.save_viz and csv_save_path is not None:
+        csv_file_path = os.path.join(csv_save_path, f"{dataset_name}_results_rank{rank}.csv")
+
+        try:
+            with open(csv_file_path, 'w', newline='') as csvfile:
+                fieldnames = ['GPU_Rank', 'Sample_ID', 'Image_Name', 'abs_rel', 'rmse', 'delta1', 'processing_time']
+                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+                writer.writeheader()
+                for row in results_data:
+                    writer.writerow(row)
+            print(f"[GPU:{rank}] 结果已保存至CSV: {csv_file_path}", file=sys.stderr)
+        except Exception as e:
+            print(f"[GPU:{rank}] 保存CSV失败: {e}", file=sys.stderr)
+
+    return metric_tracker_Lpred, metric_tracker_Rpred, processing_times
+
+
+def evaluation_normal_custom_parallel(rank, world_size, output_dir, base_data_dir, dataset_split_path, pipeline, args, eval_datasets, save_pred_vis=False):
+    """
+    支持多GPU并行的normal评估函数
+    """
+    import time
+
+    os.makedirs(output_dir, exist_ok=True)
+    device = torch.device(f"cuda:{rank}")
+
+    # 为每个数据集创建结果字典
+    all_normal_errors = {}
+    all_processing_times = {}
+    all_dataset_metrics = {}
+
+    for dataset_name, split in eval_datasets:
+        # 创建数据加载器 - 减少num_workers避免资源竞争
+        try:
+            # 创建数据集
+            from infer.dataset_normal.normal_dataloader import NormalDataset
+            dataset = NormalDataset(base_data_dir, dataset_split_path, dataset_name=dataset_name, split=split, mode='test', epoch=0)
+
+            total_samples = len(dataset)
+            if args.num_samples > 0:
+                total_samples = min(args.num_samples, total_samples)
+
+            # 计算当前GPU需要处理的样本范围
+            samples_per_gpu = total_samples // world_size
+            start_idx = rank * samples_per_gpu
+            if rank == world_size - 1:
+                end_idx = total_samples
+            else:
+                end_idx = start_idx + samples_per_gpu
+
+            # 创建样本索引并使用SubsetRandomSampler
+            from torch.utils.data import SubsetRandomSampler
+            indices = list(range(start_idx, end_idx))
+
+            dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=False, sampler=SubsetRandomSampler(indices))
+
+            if rank == 0:
+                print(f"[GPU:{rank}] 开始评估Normal数据集: {dataset_name}")
+
+        except Exception as e:
+            print(f"[GPU:{rank}] 创建数据加载器失败: {e}")
+            continue
+
+        dataset_output_dir = os.path.join(output_dir, dataset_name)
+        os.makedirs(dataset_output_dir, exist_ok=True)
+
+        if save_pred_vis:
+            save_vis_path = os.path.join(dataset_output_dir, "vis")
+            os.makedirs(save_vis_path, exist_ok=True)
+        else:
+            save_vis_path = None
+
+        processing_times = []
+        total_normal_errors = None
+        sample_count = 0
+        vis_processes = []  # 用于管理normal可视化进程
+        max_vis_processes = 5  # 限制同时运行的可视化进程数量（normal可视化更耗内存）
+
+        for data_dict in dataloader:
+            sample_count += 1
+
+            img = data_dict['img'].to(device)
+            scene_names = data_dict['scene_name']
+            img_names = data_dict['img_name']
+
+            # 获取原始图像尺寸
+            _, _, orig_H, orig_W = img.shape
+
+            start_time = time.time()
+            image_list,L_pred , norm_out = pipeline.generate_image("Predict the depth map for the image on the left and the normal map on the right.", negative_prompt="", ref_images=img, num_samples=1, num_steps=args.num_steps, cfg_guidance=args.cfg_guidance, seed=args.seed + rank, show_progress=False, size_level=args.size_level, args=args, judge=data_dict['normal'].to(device) if dataset_name == "vkitti" or dataset_name == "hypersim" else None, name=img_names)
+            end_time = time.time()
+            processing_times.append(end_time - start_time)
+
+            # 处理normal输出
+            norm_out = torch.nn.functional.interpolate(norm_out, size=(orig_H, orig_W), mode='bilinear', align_corners=False)
+            norm = torch.linalg.norm(norm_out, axis=1, keepdims=True)
+            norm[norm < 1e-9] = 1e-9
+            norm_out = norm_out / norm
+
+            pred_norm, pred_kappa = norm_out[:, :3, :, :], norm_out[:, 3:, :, :]
+            pred_kappa = None if pred_kappa.size(1) == 0 else pred_kappa
+
+            # 计算误��（如果有ground truth）
+            # if 'normal' in data_dict.keys():
+            gt_norm = data_dict['normal'].to(device)
+            gt_norm_mask = data_dict['normal_mask'].to(device)
+
+            pred_error = normal_utils.compute_normal_error(pred_norm, gt_norm)
+            if total_normal_errors is None:
+                total_normal_errors = pred_error[gt_norm_mask]
+            else:
+                total_normal_errors = torch.cat((total_normal_errors, pred_error[gt_norm_mask]), dim=0)
+
+            # 保存可视化结果（使用新的多进程方式）
+            if save_vis_path is not None:
+                # 限制同时运行的可视化进程数量
+                while len([p for p in vis_processes if p.is_alive()]) >= max_vis_processes:
+                    # 等待一些进程完成
+                    for p in vis_processes[:]:
+                        if not p.is_alive():
+                            vis_processes.remove(p)
+                    if len([p for p in vis_processes if p.is_alive()]) >= max_vis_processes:
+                        time.sleep(0.1)  # 短暂等待
+                
+                prefixs = ['%s_%s' % (i, j) for (i, j) in zip(scene_names, img_names)]
+                
+                # 预处理数据
+                img_data, pred_norm_data, pred_kappa_data, gt_norm_data, gt_norm_mask_data, pred_error_data = prepare_normal_data_for_process(
+                    img, pred_norm, pred_kappa, gt_norm, gt_norm_mask, pred_error
+                )
+                
+                if img_data is not None:  # 确保数据预处理成功
+                    # 创建子进程进行可视化保存
+                    vis_process = mp.Process(
+                        target=save_normal_visualization_worker,
+                        args=(
+                            save_vis_path,
+                            prefixs,
+                            img_data,
+                            pred_norm_data,
+                            pred_kappa_data,
+                            gt_norm_data,
+                            gt_norm_mask_data,
+                            pred_error_data,
+                            rank
+                        )
+                    )
+                    vis_process.start()
+                    vis_processes.append(vis_process)
+
+            # 输出进度信息
+            global_sample_idx = start_idx + sample_count
+            img_id = '_'.join([scene_names[0], img_names[0]])
+
+            if rank == 0 or sample_count % 10 == 0:  # 减少输出频率
+                print(f"[GPU:{rank}] | 样本:{global_sample_idx:03d} | ID:{img_id} | 时间:{processing_times[-1]:.2f}s| ", file=sys.stderr)
+
+        # 等待所有可视化进程结束
+        if save_pred_vis:
+            print(f"[GPU:{rank}] 等待Normal可视化进程完成...", file=sys.stderr)
+            for p in vis_processes:
+                p.join(timeout=60)  # normal可视化需要更长时间，设置60秒超时
+                if p.is_alive():
+                    print(f"[GPU:{rank}] Normal可视化进程超时，强制终止", file=sys.stderr)
+                    p.terminate()
+            print(f"[GPU:{rank}] 所有Normal可视化进程已完成", file=sys.stderr)
+
+        # 计算当前GPU的指标
+        metrics = None
+        if total_normal_errors is not None and len(total_normal_errors) > 0:
+            metrics = normal_utils.compute_normal_metrics(total_normal_errors)
+            if rank == 0:
+                print(f"[GPU:{rank}] 数据集 {dataset_name} 部分结果:")
+                print("mean median rmse 5 7.5 11.25 22.5 30")
+                print("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f" % (metrics['mean'], metrics['median'], metrics['rmse'], metrics['a1'], metrics['a2'], metrics['a3'], metrics['a4'], metrics['a5']))
+
+        # 存储结果
+        all_normal_errors[dataset_name] = total_normal_errors.cpu() if total_normal_errors is not None else None
+        all_processing_times[dataset_name] = processing_times
+        all_dataset_metrics[dataset_name] = metrics
+
+    return all_normal_errors, all_processing_times, all_dataset_metrics
+
+def save_normal_visualization_worker(save_vis_path, prefixs, img_data, pred_norm_data, pred_kappa_data, gt_norm_data, gt_norm_mask_data, pred_error_data, rank):
+    """
+    Normal可视化保存的工作函数，在独立进程中运行
+    Args:
+        save_vis_path: 保存路径
+        prefixs: 文件名前缀列表
+        img_data: 输入图像数据 (numpy array)
+        pred_norm_data: 预测normal数据 (numpy array)
+        pred_kappa_data: 预测kappa数据 (numpy array or None)
+        gt_norm_data: GT normal数据 (numpy array)
+        gt_norm_mask_data: GT normal掩码数据 (numpy array)
+        pred_error_data: 预测误差数据 (numpy array)
+        rank: GPU rank
+    """
+    try:
+        import infer.visualize as vis_utils
+
+        # 转换为torch tensor用于可视化函数
+        img_ts = torch.from_numpy(img_data)
+        pred_norm_ts = torch.from_numpy(pred_norm_data)
+        pred_kappa_ts = torch.from_numpy(pred_kappa_data) if pred_kappa_data is not None else None
+        gt_norm_ts = torch.from_numpy(gt_norm_data)
+        gt_norm_mask_ts = torch.from_numpy(gt_norm_mask_data)
+        pred_error_ts = torch.from_numpy(pred_error_data)
+        
+        # 使用matplotlib的非交互式后端
+        import matplotlib
+        matplotlib.use('Agg')
+        
+        # 调用可视化函数
+        vis_utils.visualize_normal(save_vis_path, prefixs, img_ts, pred_norm_ts, pred_kappa_ts, gt_norm_ts, gt_norm_mask_ts, pred_error_ts)
+        
+    except Exception as e:
+        print(f"[NORMAL-VIS-Worker-{rank}] Normal可视化保存失败: {e}", file=sys.stderr)
+
+
+def prepare_normal_data_for_process(img, pred_norm, pred_kappa, gt_norm, gt_norm_mask, pred_error):
+    """
+    预处理normal数据，转换为numpy格式供子进程使用
+    """
+    try:
+        img_data = img.cpu().numpy()
+        pred_norm_data = pred_norm.cpu().numpy()
+        pred_kappa_data = pred_kappa.cpu().numpy() if pred_kappa is not None else None
+        gt_norm_data = gt_norm.cpu().numpy()
+        gt_norm_mask_data = gt_norm_mask.cpu().numpy()
+        pred_error_data = pred_error.cpu().numpy()
+        
+        return img_data, pred_norm_data, pred_kappa_data, gt_norm_data, gt_norm_mask_data, pred_error_data
+    except Exception as e:
+        print(f"处理Normal数据失败: {e}", file=sys.stderr)
+        return None, None, None, None, None, None

FE2E/infer/sampling.py

@@ -0,0 +1,47 @@
+import math
+from collections.abc import Callable
+
+import torch
+from torch import Tensor
+
+
+def get_noise(num_samples: int, height: int, width: int, device: torch.device, dtype: torch.dtype, seed: int):
+    return torch.randn(
+        num_samples,
+        16,
+        # allow for packing
+        2 * math.ceil(height / 16),
+        2 * math.ceil(width / 16),
+        device=device,
+        dtype=dtype,
+        generator=torch.Generator(device=device).manual_seed(seed),
+    )
+
+
+def time_shift(mu: float, sigma: float, t: Tensor):
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+
+def get_lin_function(x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15) -> Callable[[float], float]:
+    m = (y2 - y1) / (x2 - x1)
+    b = y1 - m * x1
+    return lambda x: m * x + b
+
+
+def get_schedule(
+    num_steps: int,
+    image_seq_len: int,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+    shift: bool = True,
+) -> list[float]:
+    # extra step for zero
+    timesteps = torch.linspace(1, 0, num_steps + 1)
+
+    # shifting the schedule to favor high timesteps for higher signal images
+    if shift:
+        # estimate mu based on linear estimation between two points
+        mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
+        timesteps = time_shift(mu, 1.0, timesteps)
+
+    return timesteps.tolist()

FE2E/infer/seed_all.py

@@ -0,0 +1,33 @@
+# Copyright 2023 Bingxin Ke, ETH Zurich. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# --------------------------------------------------------------------------
+# If you find this code useful, we kindly ask you to cite our paper in your work.
+# Please find bibtex at: https://github.com/prs-eth/Marigold#-citation
+# More information about the method can be found at https://marigoldmonodepth.github.io
+# --------------------------------------------------------------------------
+
+
+import numpy as np
+import random
+import torch
+
+
+def seed_all(seed: int = 0):
+    """
+    Set random seeds of all components.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)

FE2E/infer/util/alignment.py

@@ -0,0 +1,88 @@
+# Author: Bingxin Ke
+# Last modified: 2024-01-11
+
+import numpy as np
+import torch
+
+
+def align_depth_least_square(
+    gt_arr: np.ndarray,
+    pred_arr: np.ndarray,
+    valid_mask_arr: np.ndarray,
+    return_scale_shift=True,
+    max_resolution=None,
+):
+    ori_shape = pred_arr.shape  # input shape
+
+    gt = gt_arr.squeeze()  # [H, W]
+    pred = pred_arr.squeeze()
+    valid_mask = valid_mask_arr.squeeze()
+
+    # Downsample
+    if max_resolution is not None:
+        scale_factor = np.min(max_resolution / np.array(ori_shape[-2:]))
+        if scale_factor < 1:
+            downscaler = torch.nn.Upsample(scale_factor=scale_factor, mode="nearest")
+            gt = downscaler(torch.as_tensor(gt).unsqueeze(0)).numpy()
+            pred = downscaler(torch.as_tensor(pred).unsqueeze(0)).numpy()
+            valid_mask = (
+                downscaler(torch.as_tensor(valid_mask).unsqueeze(0).float())
+                .bool()
+                .numpy()
+            )
+
+    assert (
+        gt.shape == pred.shape == valid_mask.shape
+    ), f"{gt.shape}, {pred.shape}, {valid_mask.shape}"
+
+    gt_masked = gt[valid_mask].reshape((-1, 1))
+    pred_masked = pred[valid_mask].reshape((-1, 1))
+
+    # numpy solver
+    _ones = np.ones_like(pred_masked)
+    A = np.concatenate([pred_masked, _ones], axis=-1)
+    X = np.linalg.lstsq(A, gt_masked, rcond=None)[0]
+    scale, shift = X
+
+    aligned_pred = pred_arr * scale + shift
+
+    # restore dimensions
+    aligned_pred = aligned_pred.reshape(ori_shape)
+
+    if return_scale_shift:
+        return aligned_pred, scale, shift
+    else:
+        return aligned_pred
+
+
+# ******************** disparity space ********************
+def depth2disparity(depth, return_mask=False):
+    if isinstance(depth, torch.Tensor):
+        disparity = torch.zeros_like(depth)
+    elif isinstance(depth, np.ndarray):
+        disparity = np.zeros_like(depth)
+    non_negtive_mask = depth > 0
+    disparity[non_negtive_mask] = 1.0 / depth[non_negtive_mask]
+    if return_mask:
+        return disparity, non_negtive_mask
+    else:
+        return disparity
+
+
+def disparity2depth(disparity, **kwargs):
+    return depth2disparity(disparity, **kwargs)
+
+# ******************** log space ********************
+def depth2log_space(depth, **kwargs):
+    if isinstance(depth, torch.Tensor):
+        log_space = torch.zeros_like(depth)
+    elif isinstance(depth, np.ndarray):
+        log_space = np.zeros_like(depth)
+    non_negtive_mask = depth > 0
+    log_space[non_negtive_mask] = np.log(depth[non_negtive_mask])
+    return log_space, non_negtive_mask
+
+def log_space2depth(log_space, **kwargs):
+    depth = np.exp(log_space)
+    return depth
+

FE2E/infer/util/metric.py

@@ -0,0 +1,151 @@
+import pandas as pd
+import torch
+
+
+# Adapted from: https://github.com/victoresque/pytorch-template/blob/master/utils/util.py
+class MetricTracker:
+    def __init__(self, *keys, writer=None):
+        self.writer = writer
+        self._data = pd.DataFrame(index=keys, columns=["total", "counts", "average"])
+        self.reset()
+
+    def reset(self):
+        for col in self._data.columns:
+            self._data[col].values[:] = 0.0
+
+    def update(self, key, value, n=1):
+        if self.writer is not None:
+            self.writer.add_scalar(key, value)
+        # 确保value是数值类型
+        value = float(value) if hasattr(value, '__float__') else value.item() if hasattr(value, 'item') else float(value)
+        self._data.at[key, "total"] += value * n
+        self._data.at[key, "counts"] += n
+        self._data.at[key, "average"] = self._data.at[key, "total"] / self._data.at[key, "counts"]
+
+    def avg(self, key):
+        return self._data.average[key]
+
+    def result(self):
+        return dict(self._data.average)
+
+def pixel_mean(pred, gt, valid_mask):
+    if valid_mask is not None:
+        masked_pred = pred * valid_mask
+        masked_gt = gt * valid_mask
+
+        valid_pixel_count = torch.sum(valid_mask, dim=(0,1))
+
+        pred_mean = torch.sum(masked_pred, dim=(0,1)) / valid_pixel_count
+        gt_mean = torch.sum(masked_gt, dim=(0,1)) / valid_pixel_count
+    else:
+        pred_mean = torch.mean(pred, dim=(0,1))
+        gt_mean = torch.mean(gt, dim=(0,1))
+
+    mean_difference = torch.abs(pred_mean - gt_mean)
+    return mean_difference
+
+def pixel_var(pred, gt, valid_mask):
+    if valid_mask is not None:
+        masked_pred = pred * valid_mask
+        masked_gt = gt * valid_mask
+
+        valid_pixel_count = torch.sum(valid_mask, dim=(0,1))
+
+        pred_mean = torch.sum(masked_pred, dim=(0,1)) / valid_pixel_count
+        gt_mean = torch.sum(masked_gt, dim=(0,1)) / valid_pixel_count
+
+        pred_var = torch.sum(valid_mask * (pred - pred_mean)**2, dim=(0,1)) / valid_pixel_count
+        gt_var = torch.sum(valid_mask * (gt - gt_mean)**2, dim=(0,1)) / valid_pixel_count
+    else:
+        pred_var = torch.var(pred, dim=(0,1))
+        gt_var = torch.var(gt, dim=(0,1))
+
+    var_difference = torch.abs(pred_var - gt_var)
+
+    return var_difference
+
+def abs_relative_difference(output, target, valid_mask=None):
+    actual_output = output
+    actual_target = target
+    abs_relative_diff = torch.abs(actual_output - actual_target) / actual_target
+    if valid_mask is not None:
+        abs_relative_diff[~valid_mask] = 0
+        n = valid_mask.sum((-1, -2))
+    else:
+        n = output.shape[-1] * output.shape[-2]
+    abs_relative_diff = torch.sum(abs_relative_diff, (-1, -2)) / n
+    return abs_relative_diff.mean()
+
+
+def squared_relative_difference(output, target, valid_mask=None):
+    actual_output = output
+    actual_target = target
+    square_relative_diff = (
+        torch.pow(torch.abs(actual_output - actual_target), 2) / actual_target
+    )
+    if valid_mask is not None:
+        square_relative_diff[~valid_mask] = 0
+        n = valid_mask.sum((-1, -2))
+    else:
+        n = output.shape[-1] * output.shape[-2]
+    square_relative_diff = torch.sum(square_relative_diff, (-1, -2)) / n
+    return square_relative_diff.mean()
+
+
+def rmse_linear(output, target, valid_mask=None):
+    actual_output = output
+    actual_target = target
+    diff = actual_output - actual_target
+    if valid_mask is not None:
+        diff[~valid_mask] = 0
+        n = valid_mask.sum((-1, -2))
+    else:
+        n = output.shape[-1] * output.shape[-2]
+    diff2 = torch.pow(diff, 2)
+    mse = torch.sum(diff2, (-1, -2)) / n
+    rmse = torch.sqrt(mse)
+    return rmse.mean()
+
+
+def rmse_log(output, target, valid_mask=None):
+    diff = torch.log(output) - torch.log(target)
+    if valid_mask is not None:
+        diff[~valid_mask] = 0
+        n = valid_mask.sum((-1, -2))
+    else:
+        n = output.shape[-1] * output.shape[-2]
+    diff2 = torch.pow(diff, 2)
+    mse = torch.sum(diff2, (-1, -2)) / n  # [B]
+    rmse = torch.sqrt(mse)
+    return rmse.mean()
+
+
+
+
+# adapt from: https://github.com/imran3180/depth-map-prediction/blob/master/main.py
+def threshold_percentage(output, target, threshold_val, valid_mask=None):
+    d1 = output / target
+    d2 = target / output
+    max_d1_d2 = torch.max(d1, d2)
+    bit_mat = (max_d1_d2 < threshold_val).to(output.dtype)
+    if valid_mask is not None:
+        bit_mat = bit_mat * valid_mask.to(output.dtype)
+        n = valid_mask.sum((-1, -2))
+    else:
+        n = torch.tensor(output.shape[-1] * output.shape[-2], device=output.device)
+    n = torch.clamp(n, min=1)
+    count_mat = torch.sum(bit_mat, (-1, -2))
+    threshold_mat = count_mat / n.to(count_mat.dtype)
+    return threshold_mat.mean()
+
+
+def delta1_acc(pred, gt, valid_mask):
+    return threshold_percentage(pred, gt, 1.25, valid_mask)
+
+
+def delta2_acc(pred, gt, valid_mask):
+    return threshold_percentage(pred, gt, 1.25**2, valid_mask)
+
+
+def delta3_acc(pred, gt, valid_mask):
+    return threshold_percentage(pred, gt, 1.25**3, valid_mask)

FE2E/infer/util/normal_utils.py

@@ -0,0 +1,78 @@
+import os
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+
+
+def get_padding(orig_H, orig_W):
+    """ returns how the input of shape (orig_H, orig_W) should be padded
+        this ensures that both H and W are divisible by 32
+    """
+    if orig_W % 32 == 0:
+        l = 0
+        r = 0
+    else:
+        new_W = 32 * ((orig_W // 32) + 1)
+        l = (new_W - orig_W) // 2
+        r = (new_W - orig_W) - l
+
+    if orig_H % 32 == 0:
+        t = 0
+        b = 0
+    else:
+        new_H = 32 * ((orig_H // 32) + 1)
+        t = (new_H - orig_H) // 2
+        b = (new_H - orig_H) - t
+    return l, r, t, b
+
+def pad_input(img, intrins, lrtb=(0,0,0,0)):
+    """ pad input image
+        img should be a torch tensor of shape (B, 3, H, W)
+        intrins should be a torch tensor of shape (B, 3, 3)
+    """
+    l, r, t, b = lrtb
+    if l+r+t+b != 0:
+        pad_value_R = (0 - 0.485) / 0.229
+        pad_value_G = (0 - 0.456) / 0.224
+        pad_value_B = (0 - 0.406) / 0.225
+
+        img_R = F.pad(img[:,0:1,:,:], (l, r, t, b), mode="constant", value=pad_value_R)
+        img_G = F.pad(img[:,1:2,:,:], (l, r, t, b), mode="constant", value=pad_value_G)
+        img_B = F.pad(img[:,2:3,:,:], (l, r, t, b), mode="constant", value=pad_value_B)
+
+        img = torch.cat([img_R, img_G, img_B], dim=1)
+
+        if intrins is not None:
+            intrins[:, 0, 2] += l
+            intrins[:, 1, 2] += t
+    return img, intrins
+
+def compute_normal_error(pred_norm, gt_norm):
+    """ compute per-pixel surface normal error in degrees
+        NOTE: pred_norm and gt_norm should be torch tensors of shape (B, 3, ...)
+    """
+    pred_error = torch.cosine_similarity(pred_norm, gt_norm, dim=1)
+    pred_error = torch.clamp(pred_error, min=-1.0, max=1.0)
+    pred_error = torch.acos(pred_error) * 180.0 / np.pi
+    pred_error = pred_error.unsqueeze(1)    # (B, 1, ...)
+    return pred_error
+
+def compute_normal_metrics(total_normal_errors):
+    """ compute surface normal metrics (used for benchmarking)
+        NOTE: total_normal_errors should be a 1D torch tensor of errors in degrees
+    """
+    total_normal_errors = total_normal_errors.detach().cpu().numpy()
+    num_pixels = total_normal_errors.shape[0]
+
+    metrics = {
+        'mean': np.average(total_normal_errors),
+        'median': np.median(total_normal_errors),
+        'rmse': np.sqrt(np.sum(total_normal_errors * total_normal_errors) / num_pixels),
+        'a1': 100.0 * (np.sum(total_normal_errors < 5) / num_pixels),
+        'a2': 100.0 * (np.sum(total_normal_errors < 7.5) / num_pixels),
+        'a3': 100.0 * (np.sum(total_normal_errors < 11.25) / num_pixels),
+        'a4': 100.0 * (np.sum(total_normal_errors < 22.5) / num_pixels),
+        'a5': 100.0 * (np.sum(total_normal_errors < 30) / num_pixels)
+    }
+    return metrics

FE2E/infer/visualize.py

@@ -0,0 +1,130 @@
+import cv2
+import numpy as np
+
+import torch
+
+from matplotlib import cm
+import matplotlib.pyplot as plt
+
+import logging
+logger = logging.getLogger('root')
+
+
+
+def tensor_to_numpy(tensor_in):
+    """ torch tensor to numpy array
+    """
+    if tensor_in is not None:
+        if tensor_in.ndim == 3:
+            # (C, H, W) -> (H, W, C)
+            tensor_in = tensor_in.detach().cpu().permute(1, 2, 0).numpy()
+        elif tensor_in.ndim == 4:
+            # (B, C, H, W) -> (B, H, W, C)
+            tensor_in = tensor_in.detach().cpu().permute(0, 2, 3, 1).numpy()
+        else:
+            raise Exception('invalid tensor size')
+    return tensor_in
+
+# def unnormalize(img_in, img_stats={'mean': [0.485, 0.456, 0.406], 
+#                                     'std': [0.229, 0.224, 0.225]}):
+def unnormalize(img_in, img_stats={'mean': [0.5,0.5,0.5], 'std': [0.5,0.5,0.5]}):
+    """ unnormalize input image
+    """
+    if torch.is_tensor(img_in):
+        img_in = tensor_to_numpy(img_in)
+
+    # 检查输入图像的数值范围，决定是否需要去归一化
+    img_min, img_max = img_in.min(), img_in.max()
+    
+    # 如果图像已经在[0,1]范围内，直接转换为[0,255]
+    if img_min >= -0.1 and img_max <= 1.1:  # 允许小的浮点误差
+        img_out = np.clip(img_in, 0, 1)
+        img_out = (img_out * 255.0).astype(np.uint8)
+    else:
+        # 如果图像在[-1,1]或其他归一化范围内，进行标准去归一化
+        img_out = np.zeros_like(img_in)
+        for ich in range(3):
+            img_out[..., ich] = img_in[..., ich] * img_stats['std'][ich]
+            img_out[..., ich] += img_stats['mean'][ich]
+        img_out = np.clip(img_out, 0, 1)
+        img_out = (img_out * 255.0).astype(np.uint8)
+    
+    return img_out
+
+def normal_to_rgb(normal, normal_mask=None):
+    """ surface normal map to RGB
+        (used for visualization)
+
+        NOTE: x, y, z are mapped to R, G, B
+        NOTE: [-1, 1] are mapped to [0, 255]
+    """
+    if torch.is_tensor(normal):
+        normal = tensor_to_numpy(normal)
+        normal_mask = tensor_to_numpy(normal_mask)
+
+    normal_norm = np.linalg.norm(normal, axis=-1, keepdims=True)
+    normal_norm[normal_norm < 1e-12] = 1e-12
+    normal = normal / normal_norm
+
+    normal_rgb = (((normal + 1) * 0.5) * 255).astype(np.uint8)
+    if normal_mask is not None:
+        normal_rgb = normal_rgb * normal_mask     # (B, H, W, 3)
+    return normal_rgb
+
+def kappa_to_alpha(pred_kappa, to_numpy=True):
+    """ Confidence kappa to uncertainty alpha
+        Assuming AngMF distribution (introduced in https://arxiv.org/abs/2109.09881)
+    """
+    if torch.is_tensor(pred_kappa) and to_numpy:
+        pred_kappa = tensor_to_numpy(pred_kappa)
+
+    if torch.is_tensor(pred_kappa):
+        alpha = ((2 * pred_kappa) / ((pred_kappa ** 2.0) + 1)) \
+            + ((torch.exp(- pred_kappa * np.pi) * np.pi) / (1 + torch.exp(- pred_kappa * np.pi)))
+        alpha = torch.rad2deg(alpha)
+    else:
+        alpha = ((2 * pred_kappa) / ((pred_kappa ** 2.0) + 1)) \
+                + ((np.exp(- pred_kappa * np.pi) * np.pi) / (1 + np.exp(- pred_kappa * np.pi)))
+        alpha = np.degrees(alpha)
+
+    return alpha
+
+
+def visualize_normal(target_dir, prefixs, img, pred_norm, pred_kappa,
+                        gt_norm, gt_norm_mask, pred_error, num_vis=-1):
+    """ visualize normal
+    """
+    error_max = 60.0
+
+    # img = tensor_to_numpy(img)                      # (B, H, W, 3)
+    pred_norm = tensor_to_numpy(pred_norm)          # (B, H, W, 3)
+    # pred_kappa = tensor_to_numpy(pred_kappa)        # (B, H, W, 1)
+    gt_norm = tensor_to_numpy(gt_norm)              # (B, H, W, 3)
+    gt_norm_mask = tensor_to_numpy(gt_norm_mask)    # (B, H, W, 1)
+    pred_error = tensor_to_numpy(pred_error)        # (B, H, W, 1)
+
+    num_vis = len(prefixs) if num_vis == -1 else num_vis
+    for i in range(num_vis):
+        # # img
+        # img_ = unnormalize(img[i, ...])
+        # target_path = '%s/%s_img.png' % (target_dir, prefixs[i])
+        # plt.imsave(target_path, img_)
+
+        # pred_norm 
+        target_path = '%s/%s_norm.png' % (target_dir, prefixs[i])
+        plt.imsave(target_path, normal_to_rgb(pred_norm[i, ...]))
+
+        # # pred_kappa
+        # if pred_kappa is not None:
+        #     pred_alpha = kappa_to_alpha(pred_kappa[i, :, :, 0])
+        #     target_path = '%s/%s_pred_alpha.png' % (target_dir, prefixs[i])
+        #     plt.imsave(target_path, pred_alpha, vmin=0.0, vmax=error_max, cmap='jet')
+
+        # gt_norm, pred_error
+        if gt_norm is not None:
+            target_path = '%s/%s_gt.png' % (target_dir, prefixs[i])
+            # plt.imsave(target_path, normal_to_rgb(gt_norm[i, ...], gt_norm_mask[i, ...]))
+
+            E = pred_error[i, :, :, 0] * gt_norm_mask[i, :, :, 0]
+            target_path = '%s/%s_pred_error.png' % (target_dir, prefixs[i])
+            plt.imsave(target_path, E, vmin=0, vmax=error_max, cmap='jet')