OASIS-DC/eval.py at main · VIPLabGachon/OASIS-DC · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
#!/usr/bin/env python3
"""
Unified evaluation script for OASIS-DC depth completion.
"""
import os
import sys
import argparse

import torch

from models.model import OASIS_DC
from train import load_config, apply_shot_adaptations, create_model_config, get_trainer


def main():
    parser = argparse.ArgumentParser(description="Evaluate OASIS-DC depth completion model")
    parser.add_argument("--dataset", type=str, required=True, choices=["nyu", "kitti"],
                        help="Dataset to use (nyu or kitti)")
    parser.add_argument("--config", type=str, required=True,
                        help="Path to configuration file")
    parser.add_argument("--checkpoint", type=str, default=None,
                        help="Path to model checkpoint (optional for Poisson-only)")
    parser.add_argument("--mode", type=str, default="all",
                        choices=["all", "full", "residual_off", "poisson_only"],
                        help="Evaluation mode")
    parser.add_argument("--device", type=str, default="cuda",
                        help="Device to use (cuda or cpu)")
    parser.add_argument("--output", type=str, default="results.json",
                        help="Output file for results")
    parser.add_argument("--save-predictions", action="store_true",
                        help="Save prediction images")
    parser.add_argument("--pred-dir", type=str, default="predictions",
                        help="Directory to save predictions")
    parser.add_argument("--shots", type=int, default=None,
                        help="Number of shots (for n-shot configs)")

    args = parser.parse_args()

    # Load configuration
    config = load_config(args.config)
    config["dataset"] = args.dataset

    # Apply shot adaptations if shots parameter is provided
    if args.shots is not None:
        config = apply_shot_adaptations(config, args.shots)
        print(f"Evaluating with {args.shots}-shot configuration")

    # Set device
    if args.device == "cuda" and not torch.cuda.is_available():
        print("CUDA not available, using CPU")
        device = torch.device("cpu")
    else:
        device = torch.device(args.device)

    print(f"Using device: {device}")
    print(f"Dataset: {args.dataset}")
    print(f"Config: {args.config}")
    print(f"Evaluation mode: {args.mode}")

    # Load checkpoint first: the model configuration it was trained with defines the architecture
    checkpoint = None
    if args.checkpoint:
        print(f"Loading checkpoint: {args.checkpoint}")
        checkpoint = torch.load(args.checkpoint, map_location=device)
        saved_model_config = (checkpoint.get("config") or {}).get("model")
        if saved_model_config:
            config["model"] = {**config.get("model", {}), **saved_model_config}
            print("Using the model configuration stored in the checkpoint")
    elif args.mode in ["full", "residual_off"]:
        print("Warning: No checkpoint provided but mode requires trained model")

    # Create model
    model_cfg = create_model_config(config)
    model = OASIS_DC(model_cfg)

    if checkpoint is not None:
        model.load_state_dict(checkpoint["model_state_dict"])
        print(f"Loaded checkpoint from epoch {checkpoint.get('epoch', 'unknown')}")

    # Get trainer
    trainer = get_trainer(args.dataset, model, device, config)

    # Get test dataloader
    test_loader = trainer.get_dataloader("test")
    print(f"Test samples: {len(test_loader.dataset)}")

    # Run evaluation
    print("\n=== Starting Evaluation ===")

    if args.mode == "all":
        # Evaluate all modes
        results = trainer.evaluate_all_modes(test_loader)

    else:
        # Evaluate specific mode
        results = trainer.evaluate(test_loader, mode=args.mode)

    # Print results
    trainer.print_results(results)

    # Save results
    print(f"\nSaving results to: {args.output}")
    trainer.save_results(results, args.output)

    # Save predictions if requested
    if args.save_predictions:
        print(f"Prediction saving not yet implemented")

    print("\nEvaluation completed!")


if __name__ == "__main__":
    main()