Cannot Reproduce Evaluation Result on ETH-3D Dataset

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from statistics import mean

import numpy as np
import torch
import torchvision.transforms as T
from tqdm import tqdm

from transformers import DepthAnythingForDepthEstimation

from Utility.HDF5 import H5Store
from .Dataset import All_ETH3D_Data


@torch.no_grad()
def align_depth(gt: torch.Tensor, pred: torch.Tensor, eps: float = 1e-6):
    """
    Aligns pred to gt using inverse-depth linear scaling:
      inv_gt ≈ scale * inv_pred + shift
    where inv_* = 1 / (depth + eps).
    
    Args:
        gt:   [1,1,H,W] sparse ground-truth depth (torch.Tensor)
        pred: [1,1,H,W] dense predicted depth (torch.Tensor)
        eps:  small constant to avoid div by zero
        
    Returns:
        aligned: [1,1,H,W] aligned depth map
        scale:   scalar scale factor (torch.Tensor)
        shift:   scalar shift term (torch.Tensor)
    """
    # flatten valid pixels
    mask       = torch.isfinite(gt)
    gt_vals    = gt[mask]
    pred_vals  = pred[mask]
    
    # inverse-depth
    inv_gt     = 1.0 / (gt_vals + eps)
    inv_pred   = 1.0 / (pred_vals + eps)
    
    # build normal equations A^T A x = A^T b
    A          = torch.stack([inv_pred, torch.ones_like(inv_pred)], dim=1)  # [N,2]
    ATA        = A.T @ A                                                   # [2,2]
    ATb        = A.T @ inv_gt                                               # [2]
    
    # solve for [scale, shift]
    x          = torch.linalg.solve(ATA, ATb)
    scale, shift = x[0], x[1]
    
    # apply to full pred map
    inv_pred_full = 1.0 / (pred + eps)
    aligned       = 1.0 / (inv_pred_full * scale + shift + eps)
    
    return aligned, scale, shift

@torch.no_grad()
def reshape_input(image: torch.Tensor, depth: torch.Tensor, target_size: tuple[int, int]):
    image_resized = torch.nn.functional.interpolate(image, size=target_size, mode="bilinear")
    depth_resized = torch.nn.functional.interpolate(depth, size=target_size, mode="nearest")
    return image_resized, depth_resized


def infer_once(variant: str, img_size: tuple[int, int]):
    # Paths & constants
    RESULT_ROOT = Path(f"Result/ETH3D/")
    H5_PATH     = RESULT_ROOT / f"DA2_{variant}_size_{img_size[0]}_{img_size[1]}.h5"
    METRIC_PATH = RESULT_ROOT / f"DA2_{variant}_size_{img_size[0]}_{img_size[1]}_metric.json"
    DEVICE      = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print(f"Loading DINOv2‑{variant} backbone and DPT head …")

    model = DepthAnythingForDepthEstimation\
        .from_pretrained(f"depth-anything/Depth-Anything-V2-{variant.capitalize()}-hf")\
        .to(DEVICE)\
        .eval()

    preprocess = T.Compose([
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    print("Building dataset & HDF5 store …")

    schema = {
        "scaled_gt_depth"   : (np.float32, img_size),
        "orig_pred_depth"   : (np.float32, img_size),
        "aligned_pred_depth": (np.float32, img_size),
        "scale_factor"      : (np.float32, ()),
        "absrel"            : (np.float32, ()),
        "l1"                : (np.float32, ()),
        "infer_time_sec"    : (np.float32, ()),
    }
    absrel_arr: list[float] = []
    l1_arr    : list[float] = []

    with H5Store(H5_PATH, mode="w", schema=schema, chunk_size=1) as store, \
         tqdm(range(len(All_ETH3D_Data))) as pb:
        for idx in pb:
            image, depth = reshape_input(*All_ETH3D_Data[idx], target_size=img_size)

            with torch.no_grad():
                image = preprocess(image).to(DEVICE)
                
                torch.cuda.synchronize()
                start_time = time.time()
                
                pred = model(image).predicted_depth
                
                torch.cuda.synchronize()
                end_time = time.time()
        
                pred = (torch.nn.functional.interpolate(
                    pred.unsqueeze(0), img_size, mode="bilinear", align_corners=True
                )).cpu()
                gt   = depth
            
            valid_mask = gt.isfinite() & (gt > 0)
            
            # scale      = (gt[valid_mask].mean() / pred[valid_mask].mean()).item()
            # aligned    = pred * scale
            aligned, scale, shift = align_depth(gt, pred)
            
            
            absrel = torch.mean((aligned[valid_mask] - gt[valid_mask]).abs() / gt[valid_mask]).item()
            l1   = (aligned[valid_mask] - gt[valid_mask]).abs().mean().item()
            dt   = end_time - start_time
            
            l1_arr  .append(l1)
            absrel_arr.append(absrel)

            store.push({
                "scaled_gt_depth"   : np.expand_dims(gt.squeeze().cpu().numpy().astype(np.float32), 0),
                "orig_pred_depth"   : np.expand_dims(pred.squeeze().cpu().numpy().astype(np.float32), 0),
                "aligned_pred_depth": np.expand_dims(aligned.squeeze().cpu().numpy().astype(np.float32), 0),
                "scale_factor"      : np.array([scale], np.float32),
                "absrel"            : np.array([absrel],  np.float32),
                "l1"                : np.array([l1],    np.float32),
                "infer_time_sec"    : np.array([dt],    np.float32),
            })
            
            pb.set_postfix(l1=l1, absrel=absrel, t_sec=dt)
                
        summary = {"absrel": absrel_arr, "l1": l1_arr}
        with open(METRIC_PATH, "w") as f:
            json.dump(summary, f)

        print("\nFinished.")
        print(f"Average AbsRel: {mean(absrel_arr):.3f}")
        print(f"Average L1  : {mean(l1_arr):.3f}")
        print(f"Results written to {H5_PATH}")



if __name__ == "__main__":
    # CLI
    parser = argparse.ArgumentParser(description="Predict depth on ETH-3D with DepthAnything-v2")
    subparser = parser.add_subparsers(dest="action", required=True)
    
    # Single Run
    d_parser = subparser.add_parser("direct")
    d_parser.add_argument("variant", choices=["small", "large"], help="DepthAnything model size")
    d_parser.add_argument("--size", type=int, nargs=2, default=(2016, 3024))
    
    args = parser.parse_args()

    match args.action:
        case "direct":
            variant : str = args.variant
            img_size: tuple[int, int] = args.size
            infer_once(variant, img_size)