File size: 4,098 Bytes
98a3af2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
114
115
116
117
118
119
120
#!/usr/bin/env python3
"""
DEIM Dataset Visualization Entry Point
Copyright (c) 2024 The DEIM Authors. All Rights Reserved.

Simple script to launch the dataset visualizer for any configured dataset.
"""

import os
import sys
import argparse
from pathlib import Path

# Add project root to Python path
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from tools.visualization.dataset_visualizer import DatasetVisualizer


def list_available_configs():
    """List all available dataset configuration files."""
    config_dir = project_root / "configs" / "dataset"
    if not config_dir.exists():
        print("No dataset configurations found.")
        return []
    
    configs = list(config_dir.glob("*.yml"))
    print("\nAvailable dataset configurations:")
    for i, config in enumerate(configs, 1):
        print(f"  {i:2d}. {config.stem}")
    
    return configs


def main():
    """Main entry point for dataset visualization."""
    parser = argparse.ArgumentParser(
        description='DEIM Dataset Visualizer - Interactive dataset browser with OpenCV',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python visualize_dataset.py coco_detection          # Visualize COCO dataset (train split)
  python visualize_dataset.py coco_detection --split val  # Visualize COCO dataset (val split)
  python visualize_dataset.py --list                  # List all available configurations
  python visualize_dataset.py drone_detection --start-index 100  # Start from sample 100
  
  # Or use full path to config file:
  python visualize_dataset.py configs/dataset/custom_detection.yml
        """
    )
    
    parser.add_argument('config', nargs='?', help='Dataset configuration name or path to YAML file')
    parser.add_argument('--split', choices=['train', 'val'], default='train',
                        help='Dataset split to visualize (default: train)')
    parser.add_argument('--start-index', type=int, default=0,
                        help='Starting index for visualization (default: 0)')
    parser.add_argument('--list', action='store_true',
                        help='List available dataset configurations')
    
    args = parser.parse_args()
    
    # Handle list option
    if args.list:
        list_available_configs()
        return
    
    # Validate config argument
    if not args.config:
        print("Error: Please specify a dataset configuration.")
        print("Use --list to see available configurations.")
        return
    
    # Determine config file path
    if args.config.endswith('.yml') or args.config.endswith('.yaml'):
        # Full path provided
        config_path = args.config
    else:
        # Configuration name provided
        config_path = project_root / "configs" / "dataset" / f"{args.config}.yml"
    
    # Validate config file exists
    if not os.path.exists(config_path):
        print(f"Error: Configuration file '{config_path}' not found")
        print("\nAvailable configurations:")
        list_available_configs()
        return
    
    # Print information
    print("=" * 60)
    print("DEIM Dataset Visualizer")
    print("=" * 60)
    print(f"Configuration: {config_path}")
    print(f"Split: {args.split}")
    print(f"Starting index: {args.start_index}")
    print()
    
    try:
        # Create and run visualizer
        visualizer = DatasetVisualizer(str(config_path), args.split)
        
        # Validate starting index
        if args.start_index >= len(visualizer.dataset):
            print(f"Warning: Start index {args.start_index} is out of range. Using 0.")
            args.start_index = 0
        
        visualizer.current_index = args.start_index
        visualizer.run()
        
    except KeyboardInterrupt:
        print("\nVisualization interrupted by user.")
    except Exception as e:
        print(f"Error: {e}")
        print("\nTip: Make sure the dataset paths in the configuration file are correct")
        print("     and that all required dependencies are installed.")


if __name__ == '__main__':
    main()