File size: 1,600 Bytes
a0c7025 a33e03b a0c7025 3441a79 2dd2ae5 c653e7e d58a9b6 1fe2937 a0c7025 d58a9b6 1fe2937 97e9dcb f5518c0 c653e7e 7d7e199 f5518c0 1fe2937 2dd2ae5 a0c7025 1fe2937 a0c7025 d58a9b6 a0c7025 3441a79 a0c7025 |
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 |
import sys
from pathlib import Path
import hydra
project_root = Path(__file__).resolve().parent.parent
sys.path.append(str(project_root))
from yolo.config.config import Config
from yolo.model.yolo import create_model
from yolo.tools.data_loader import create_dataloader
from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
from yolo.utils.bounding_box_utils import Vec2Box
from yolo.utils.deploy_utils import FastModelLoader
from yolo.utils.logging_utils import ProgressLogger
from yolo.utils.model_utils import get_device
@hydra.main(config_path="config", config_name="config", version_base=None)
def main(cfg: Config):
progress = ProgressLogger(cfg, exp_name=cfg.name)
device, use_ddp = get_device(cfg.device)
dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task, use_ddp)
if getattr(cfg.task, "fast_inference", False):
model = FastModelLoader(cfg).load_model(device)
else:
model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight)
model = model.to(device)
vec2box = Vec2Box(model, cfg.image_size, device)
if cfg.task.task == "train":
trainer = ModelTrainer(cfg, model, vec2box, progress, device, use_ddp)
trainer.solve(dataloader)
if cfg.task.task == "inference":
tester = ModelTester(cfg, model, vec2box, progress, device)
tester.solve(dataloader)
if cfg.task.task == "validation":
valider = ModelValidator(cfg.task, model, vec2box, progress, device)
valider.solve(dataloader)
if __name__ == "__main__":
main()
|