mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-03-12 12:25:32 +08:00
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
import argparse
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"action",
|
|
choices=["train", "benchmark", "visualize", "generate"],
|
|
help="Action to perform: train, benchmark, visualize, or generate",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if args.action == "train":
|
|
from compressors import train
|
|
|
|
train(
|
|
epoch_size=10, batch_size=64, lr=1e-4, checkpoint_path="hash_checkpoint.pt"
|
|
)
|
|
elif args.action == "benchmark":
|
|
from typing import cast
|
|
|
|
import torch
|
|
from benchmarks import run_benchmark
|
|
from compressors import DinoCompressor
|
|
from configs import cfg_manager
|
|
from transformers import AutoImageProcessor, BitImageProcessorFast
|
|
from utils import get_device
|
|
|
|
config = cfg_manager.get()
|
|
benchmark_cfg = config.benchmark
|
|
|
|
if not benchmark_cfg.enabled:
|
|
print("Benchmark is not enabled. Set benchmark.enabled=true in config.yaml")
|
|
exit(1)
|
|
|
|
device = get_device()
|
|
|
|
# Load model and processor based on config
|
|
model_cfg = config.model
|
|
processor = cast(
|
|
BitImageProcessorFast,
|
|
AutoImageProcessor.from_pretrained(model_cfg.dino_model, device_map=device),
|
|
)
|
|
|
|
# Load compressor weights if specified in model config
|
|
model = DinoCompressor().to(device)
|
|
if model_cfg.compressor_path is not None:
|
|
from compressors import HashCompressor
|
|
|
|
compressor = HashCompressor(
|
|
input_dim=model_cfg.compression_dim,
|
|
output_dim=model_cfg.compression_dim,
|
|
)
|
|
compressor.load_state_dict(torch.load(model_cfg.compressor_path))
|
|
# Wrap with compressor if path is specified
|
|
model.compressor = compressor
|
|
|
|
# Run benchmark
|
|
run_benchmark(
|
|
model=model,
|
|
processor=processor,
|
|
config=benchmark_cfg,
|
|
model_name="dinov2",
|
|
)
|
|
elif args.action == "visualize":
|
|
from visualizer import app
|
|
|
|
app.run(debug=True)
|
|
else: # generate
|
|
from configs import cfg_manager
|
|
from data_loading.synthesizer import ImageSynthesizer
|
|
|
|
config = cfg_manager.get()
|
|
dataset_cfg = config.dataset
|
|
|
|
synthesizer = ImageSynthesizer(
|
|
dataset_root=dataset_cfg.dataset_root,
|
|
output_dir=dataset_cfg.output_dir,
|
|
num_objects_range=dataset_cfg.num_objects_range,
|
|
num_scenes=dataset_cfg.num_scenes,
|
|
object_scale_range=dataset_cfg.object_scale_range,
|
|
rotation_range=dataset_cfg.rotation_range,
|
|
overlap_threshold=dataset_cfg.overlap_threshold,
|
|
seed=dataset_cfg.seed,
|
|
)
|
|
|
|
generated_files = synthesizer.generate()
|
|
print(
|
|
f"Generated {len(generated_files)} synthesized images in {dataset_cfg.output_dir}"
|
|
)
|