|
|
from collections import namedtuple |
|
|
|
|
|
import numpy as np |
|
|
import torch |
|
|
|
|
|
from .detectors import build_detector |
|
|
|
|
|
try: |
|
|
import kornia |
|
|
except: |
|
|
pass |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def build_network(model_cfg, num_class, dataset): |
|
|
model = build_detector( |
|
|
model_cfg=model_cfg, num_class=num_class, dataset=dataset |
|
|
) |
|
|
return model |
|
|
|
|
|
|
|
|
def load_data_to_gpu(batch_dict): |
|
|
for key, val in batch_dict.items(): |
|
|
if key == 'camera_imgs': |
|
|
batch_dict[key] = val.cuda() |
|
|
elif not isinstance(val, np.ndarray): |
|
|
continue |
|
|
elif key in ['frame_id', 'metadata', 'calib', 'image_paths','ori_shape','img_process_infos']: |
|
|
continue |
|
|
elif key in ['images']: |
|
|
batch_dict[key] = kornia.image_to_tensor(val).float().cuda().contiguous() |
|
|
elif key in ['image_shape']: |
|
|
batch_dict[key] = torch.from_numpy(val).int().cuda() |
|
|
else: |
|
|
batch_dict[key] = torch.from_numpy(val).float().cuda() |
|
|
|
|
|
|
|
|
def model_fn_decorator(): |
|
|
ModelReturn = namedtuple('ModelReturn', ['loss', 'tb_dict', 'disp_dict']) |
|
|
|
|
|
def model_func(model, batch_dict): |
|
|
load_data_to_gpu(batch_dict) |
|
|
ret_dict, tb_dict, disp_dict = model(batch_dict) |
|
|
|
|
|
loss = ret_dict['loss'].mean() |
|
|
if hasattr(model, 'update_global_step'): |
|
|
model.update_global_step() |
|
|
else: |
|
|
model.module.update_global_step() |
|
|
|
|
|
return ModelReturn(loss, tb_dict, disp_dict) |
|
|
|
|
|
return model_func |
|
|
|
