def main():
global predictor
global conn
server = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
server.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)
server.bind(('127.0.0.1',8000))
server.listen(5)
print("waiting msg ...")
conn, clint_add = server.accept()
args = parse_args()
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
load_config(cfg, args.config)
logger = Logger(-1, use_tensorboard=False)
predictor = Predictor(cfg, args.model, logger, device='cuda:0')
logger.log('Press "Esc", "q" or "Q" to exit.')
if args.demo == 'image':
if os.path.isdir(args.path):
files = get_image_list(args.path)
else:
files = [args.path]
files.sort()
for image_name in files:
meta, res = predictor.inference(image_name)
predictor.visualize(res, meta, cfg.class_names, 0.35)
ch = cv2.waitKey(0)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
elif args.demo == 'video' or args.demo == 'webcam':
rospy.init_node('listener', anonymous=True)
detime = time.time()
rospy.Subscriber("/camera/color/image_raw", Image, callback)
rospy.spin()
def main():
args = parse_args()
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
load_config(cfg, args.config)
logger = Logger(-1, use_tensorboard=False)
predictor = Predictor(cfg, args.model, logger, device='cuda:0')
logger.log('Press "Esc", "q" or "Q" to exit.')
if args.demo == 'image':
if os.path.isdir(args.path):
files = get_image_list(args.path)
else:
files = [args.path]
files.sort()
for image_name in files:
meta, res = predictor.inference(image_name)
predictor.visualize(res, meta, cfg.class_names, 0.35)
ch = cv2.waitKey(0)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
elif args.demo == 'video' or args.demo == 'webcam':
cap = cv2.VideoCapture(args.path if args.demo ==
'video' else args.camid)
while True:
ret_val, frame = cap.read()
meta, res = predictor.inference(frame)
predictor.visualize(res, meta, cfg.class_names, 0.35)
ch = cv2.waitKey(1)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
def main():
args = parse_args()
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
load_config(cfg, args.config)
logger = Logger(-1, use_tensorboard=False)
predictor = Predictor(cfg, args.model, logger, device='cuda:0')
logger.log('Press "Esc", "q" or "Q" to exit.')
fpsReport = 0
font = cv2.FONT_HERSHEY_SIMPLEX
timeStamp = time.time()
if args.demo == 'image':
if os.path.isdir(args.path):
files = get_image_list(args.path)
else:
files = [args.path]
files.sort()
for image_name in files:
meta, res = predictor.inference(image_name)
predictor.visualize(res, meta, cfg.class_names, 0.35)
ch = cv2.waitKey(0)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
elif args.demo == 'video' or args.demo == 'webcam':
cap = cv2.VideoCapture(args.path if args.demo ==
'video' else args.camid)
while True:
ret_val, frame = cap.read()
meta, res = predictor.inference(frame)
predictor.visualize(res, meta, cfg.class_names, 0.35)
dt = time.time() - timeStamp
fps = 1 / dt
fpsReport = .9 * fpsReport + .1 * fps
print("FPS: " + str(fpsReport))
timeStamp = time.time()
ch = cv2.waitKey(1)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
elif args.demo == 'csicam':
cap = cv2.VideoCapture(gstreamer_pipeline(), cv2.CAP_GSTREAMER)
while True:
ret_val, frame = cap.read()
meta, res = predictor.inference(frame)
predictor.visualize(res, meta, cfg.class_names, 0.35)
dt = time.time() - timeStamp
fps = 1 / dt
fpsReport = .9 * fpsReport + .1 * fps
print("FPS: " + str(fpsReport))
timeStamp = time.time()
ch = cv2.waitKey(1)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
def run():
args = parse_args()
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
# ----- load and parse config file
load_config(cfg, args.config)
# ----- set logger
logger = Logger(-1, use_tensorboard=False)
# ----- set device
device = torch_utils.select_device(args.device,
apex=False,
batch_size=None)
# ----- set predictor
predictor = Predictor(cfg, args.model, logger, device=device) # 'cuda:0'
logger.log('Press "Esc", "q" or "Q" to exit.')
if args.demo == 'image':
if os.path.isdir(args.path):
files = get_image_list(args.path)
else:
files = [args.path]
files.sort()
for img_path in files:
meta, res_dict = predictor.inference(img_path)
predictor.visualize(img_path, res_dict, meta, cfg.class_names,
0.35)
ch = cv2.waitKey(0)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
elif args.demo == 'video' or args.demo == 'webcam':
cap = cv2.VideoCapture(args.path if args.demo ==
'video' else args.camid)
while True:
ret_val, frame = cap.read()
# ----- inference
meta, res_dict = predictor.inference(frame)
# -----
predictor.visualize(res_dict, meta, cfg.class_names, 0.35)
ch = cv2.waitKey(1)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
def main(args):
warnings.warn('Warning! Old training code is deprecated and will be deleted '
'in next version. Please use tools/train.py')
load_config(cfg, args.config)
local_rank = int(args.local_rank)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
mkdir(local_rank, cfg.save_dir) # mkdir用@rank_filter包裹，主进程创建save_dir
logger = Logger(local_rank, cfg.save_dir)
if args.seed is not None:
logger.log('Set random seed to {}'.format(args.seed))
init_seeds(args.seed)
logger.log('Creating model...')
model = build_model(cfg.model)
logger.log('Setting up data...')
train_dataset = build_dataset(cfg.data.train, 'train')
val_dataset = build_dataset(cfg.data.val, 'test')
if len(cfg.device.gpu_ids) > 1:
print('rank = ', local_rank)
num_gpus = torch.cuda.device_count()
torch.cuda.set_device(local_rank % num_gpus)
dist.init_process_group(backend='nccl')
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=cfg.device.batchsize_per_gpu,
num_workers=cfg.device.workers_per_gpu, pin_memory=True,
collate_fn=collate_function, sampler=train_sampler,
drop_last=True)
else:
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=cfg.device.batchsize_per_gpu,
shuffle=True, num_workers=cfg.device.workers_per_gpu,
pin_memory=True, collate_fn=collate_function, drop_last=True)
val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=cfg.device.batchsize_per_gpu,
shuffle=False, num_workers=cfg.device.workers_per_gpu,
pin_memory=True, collate_fn=collate_function, drop_last=True)
trainer = build_trainer(local_rank, cfg, model, logger)
if 'load_model' in cfg.schedule:
trainer.load_model(cfg)
if 'resume' in cfg.schedule:
trainer.resume(cfg)
evaluator = build_evaluator(cfg, val_dataset)
logger.log('Starting training...')
trainer.run(train_dataloader, val_dataloader, evaluator)
def main():
args = parse_args()
local_rank = 0
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
load_config(cfg, args.config)
logger = Logger(local_rank, use_tensorboard=False)
predictor = Predictor(cfg, args.model, logger, device='cuda:0')
logger.log('Press "Esc", "q" or "Q" to exit.')
current_time = time.localtime()
if args.demo == 'image':
if os.path.isdir(args.path):
files = get_image_list(args.path)
else:
files = [args.path]
files.sort()
for image_name in files:
meta, res = predictor.inference(image_name)
result_image = predictor.visualize(res, meta, cfg.class_names, 0.35)
if args.save_result:
save_folder = os.path.join(cfg.save_dir, time.strftime("%Y_%m_%d_%H_%M_%S", current_time))
mkdir(local_rank, save_folder)
save_file_name = os.path.join(save_folder, os.path.basename(image_name))
cv2.imwrite(save_file_name, result_image)
ch = cv2.waitKey(0)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
elif args.demo == 'video' or args.demo == 'webcam':
cap = cv2.VideoCapture(args.path if args.demo == 'video' else args.camid)
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) # float
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) # float
fps = cap.get(cv2.CAP_PROP_FPS)
save_folder = os.path.join(cfg.save_dir, time.strftime("%Y_%m_%d_%H_%M_%S", current_time))
mkdir(local_rank, save_folder)
save_path = os.path.join(save_folder, args.path.split('/')[-1]) if args.demo == 'video' else os.path.join(save_folder, 'camera.mp4')
print(f'save_path is {save_path}')
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (int(width), int(height)))
while True:
ret_val, frame = cap.read()
if ret_val:
meta, res = predictor.inference(frame)
result_frame = predictor.visualize(res, meta, cfg.class_names, 0.35)
if args.save_result:
vid_writer.write(result_frame)
ch = cv2.waitKey(1)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
else:
break
def main(config, model_path: str, output_path: str, input_shape=(320, 320)):
logger = Logger(local_rank=-1,
save_dir=config.save_dir,
use_tensorboard=False)
# Create model and load weights
model = build_model(config.model)
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
load_model_weight(model, checkpoint, logger)
# Convert backbone weights for RepVGG models
if config.model.arch.backbone.name == "RepVGG":
deploy_config = config.model
deploy_config.arch.backbone.update({"deploy": True})
deploy_model = build_model(deploy_config)
from nanodet.model.backbone.repvgg import repvgg_det_model_convert
model = repvgg_det_model_convert(model, deploy_model)
# TorchScript: tracing the model with dummy inputs
with torch.no_grad():
dummy_input = torch.zeros(1, 3, input_shape[0],
input_shape[1]) # Batch size = 1
model.eval().cpu()
model_traced = torch.jit.trace(model,
example_inputs=dummy_input).eval()
model_traced.save(output_path)
print("Finished export to TorchScript")
def main(config, model_path, output_path, input_shape=(320, 320)):
logger = Logger(-1, config.save_dir, False)
model = build_model(config.model)
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
load_model_weight(model, checkpoint, logger)
dummy_input = torch.autograd.Variable(torch.randn(1, 3, input_shape[0], input_shape[1]))
torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=11)
print('finished exporting onnx ')
def main(config, model_path, output_path, input_shape=(320, 320)):
logger = Logger(-1, config.save_dir, False)
model = build_model(config.model)
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
load_model_weight(model, checkpoint, logger)
if config.model.arch.backbone.name == 'RepVGG':
deploy_config = config.model
deploy_config.arch.backbone.update({'deploy': True})
deploy_model = build_model(deploy_config)
from nanodet.model.backbone.repvgg import repvgg_det_model_convert
model = repvgg_det_model_convert(model, deploy_model)
dummy_input = torch.autograd.Variable(
torch.randn(1, 3, input_shape[0], input_shape[1]))
dynamic_axes = {
"input": {
0: "batch_size"
},
"output1": {
0: "batch_size"
},
"output2": {
0: "batch_size"
},
"output3": {
0: "batch_size"
},
"output4": {
0: "batch_size"
},
"output5": {
0: "batch_size"
},
"output6": {
0: "batch_size"
}
}
input_names = ['input']
output_names = [
'output1', 'output2', 'output3', 'output4', 'output5', 'output6'
]
torch.onnx.export(model,
dummy_input,
output_path,
verbose=True,
keep_initializers_as_inputs=True,
opset_version=12,
input_names=input_names,
output_names=output_names)
import onnx
from onnxsim import simplify
model = onnx.load(output_path)
# convert model
model_simp, check = simplify(model)
onnx.save(model_simp, output_path)
print('finished exporting onnx ')
def main():
args = parse_args()
if args.device != 'cpu':
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
load_config(cfg, args.config)
logger = Logger(-1, use_tensorboard=False)
predictor = Predictor(cfg, args.model, logger, device=args.device)
logger.log('Press "Esc", "q" or "Q" to exit.')
if args.demo == 'image':
if os.path.isdir(args.path):
files = get_image_list(args.path)
else:
files = [args.path]
files.sort()
for image_name in files:
meta, res = predictor.inference(image_name)
predictor.visualize(res, meta, cfg.class_names, 0.35)
ch = cv2.waitKey(0)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
elif args.demo == 'video' or args.demo == 'webcam':
cap = cv2.VideoCapture(args.path if args.demo ==
'video' else args.camid)
result_video_path = args.path.replace('.avi', '_result.avi')
fps = cap.get(cv2.CAP_PROP_FPS)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
size = (w, h)
result_cap = cv2.VideoWriter(result_video_path,
cv2.VideoWriter_fourcc(*'mp4v'), fps,
size)
while True:
ret_val, frame = cap.read()
if not ret_val:
break
meta, res = predictor.inference(frame)
# predictor.visualize(res, meta, cfg.class_names, 0.35)
result_frame = predictor.draw_bbox(res, meta, cfg.class_names,
0.36)
write_frame = cv2.resize(result_frame, (w, h),
interpolation=cv2.INTER_NEAREST)
result_cap.write(write_frame)
def main(args):
load_config(cfg, args.config)
if cfg.model.arch.head.num_classes != len(cfg.class_names):
raise ValueError('cfg.model.arch.head.num_classes must equal len(cfg.class_names),but got {} and {}'.format(cfg.model.arch.head.num_classes,len(cfg.class_names)))
local_rank = int(args.local_rank)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
mkdir(local_rank, cfg.save_dir)
logger = Logger(local_rank, cfg.save_dir)
if args.seed is not None:
logger.log('Set random seed to {}'.format(args.seed))
pl.seed_everything(args.seed)
logger.log('Setting up data...')
train_dataset = build_dataset(cfg.data.train, 'train')
val_dataset = build_dataset(cfg.data.val, 'test')
evaluator = build_evaluator(cfg, val_dataset)
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=cfg.device.batchsize_per_gpu,
shuffle=True, num_workers=cfg.device.workers_per_gpu,
pin_memory=True, collate_fn=collate_function, drop_last=True)
# TODO: batch eval
val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False,
num_workers=cfg.device.workers_per_gpu,
pin_memory=True, collate_fn=collate_function, drop_last=True)
logger.log('Creating model...')
task = TrainingTask(cfg, evaluator)
if 'load_model' in cfg.schedule:
ckpt = torch.load(cfg.schedule.load_model)
if 'pytorch-lightning_version' not in ckpt:
warnings.warn('Warning! Old .pth checkpoint is deprecated. '
'Convert the checkpoint with tools/convert_old_checkpoint.py ')
ckpt = convert_old_model(ckpt)
task.load_state_dict(ckpt['state_dict'], strict=False)
model_resume_path = os.path.join(cfg.save_dir, 'model_last.ckpt') if 'resume' in cfg.schedule else None
trainer = pl.Trainer(default_root_dir=cfg.save_dir,
max_epochs=cfg.schedule.total_epochs,
gpus=cfg.device.gpu_ids,
check_val_every_n_epoch=cfg.schedule.val_intervals,
accelerator='ddp',
log_every_n_steps=cfg.log.interval,
num_sanity_val_steps=0,
resume_from_checkpoint=model_resume_path,
callbacks=[ProgressBar(refresh_rate=0)] # disable tqdm bar
)
trainer.fit(task, train_dataloader, val_dataloader)
def main(args):
warnings.warn(
'Warning! Old testing code is deprecated and will be deleted '
'in next version. Please use tools/test.py')
load_config(cfg, args.config)
local_rank = -1
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
cfg.defrost()
timestr = datetime.datetime.now().__format__('%Y%m%d%H%M%S')
cfg.save_dir = os.path.join(cfg.save_dir, timestr)
cfg.freeze()
mkdir(local_rank, cfg.save_dir)
logger = Logger(local_rank, cfg.save_dir)
logger.log('Creating model...')
model = build_model(cfg.model)
logger.log('Setting up data...')
val_dataset = build_dataset(cfg.data.val, args.task)
val_dataloader = torch.utils.data.DataLoader(
val_dataset,
batch_size=cfg.device.batchsize_per_gpu,
shuffle=False,
num_workers=cfg.device.workers_per_gpu,
pin_memory=True,
collate_fn=collate_function,
drop_last=True)
trainer = build_trainer(local_rank, cfg, model, logger)
cfg.schedule.update({'load_model': args.model})
trainer.load_model(cfg)
evaluator = build_evaluator(cfg, val_dataset)
logger.log('Starting testing...')
with torch.no_grad():
results, val_loss_dict = trainer.run_epoch(0,
val_dataloader,
mode=args.task)
if args.task == 'test':
res_json = evaluator.results2json(results)
json_path = os.path.join(cfg.save_dir,
'results{}.json'.format(timestr))
json.dump(res_json, open(json_path, 'w'))
elif args.task == 'val':
eval_results = evaluator.evaluate(results,
cfg.save_dir,
rank=local_rank)
if args.save_result:
txt_path = os.path.join(cfg.save_dir,
"eval_results{}.txt".format(timestr))
with open(txt_path, "a") as f:
for k, v in eval_results.items():
f.write("{}: {}\n".format(k, v))
def main(config, model_path, output_path, input_shape=(320, 320)):
logger = Logger(-1, config.save_dir, False)
model = build_model(config.model)
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
load_model_weight(model, checkpoint, logger)
if config.model.arch.backbone.name == 'RepVGG':
deploy_config = config.model
deploy_config.arch.backbone.update({'deploy': True})
deploy_model = build_model(deploy_config)
from nanodet.model.backbone.repvgg import repvgg_det_model_convert
model = repvgg_det_model_convert(model, deploy_model)
dummy_input = torch.autograd.Variable(torch.randn(1, 3, input_shape[0], input_shape[1]))
torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=11)
print('finished exporting onnx ')
def main(args):
load_config(cfg, args.config)
local_rank = int(args.local_rank)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
mkdir(local_rank, cfg.save_dir)
logger = Logger(local_rank, cfg.save_dir)
# TODO: replace with lightning random seed
if args.seed is not None:
logger.log('Set random seed to {}'.format(args.seed))
init_seeds(args.seed)
logger.log('Setting up data...')
train_dataset = build_dataset(cfg.data.train, 'train')
val_dataset = build_dataset(cfg.data.val, 'test')
evaluator = build_evaluator(cfg, val_dataset)
logger.log('Creating model...')
task = TrainingTask(cfg, evaluator, logger)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.device.batchsize_per_gpu,
shuffle=True,
num_workers=cfg.device.workers_per_gpu,
pin_memory=True,
collate_fn=collate_function,
drop_last=True)
# TODO: batch eval
val_dataloader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True,
collate_fn=collate_function,
drop_last=True)
trainer = pl.Trainer(default_root_dir=cfg.save_dir,
max_epochs=cfg.schedule.total_epochs,
gpus=cfg.device.gpu_ids,
check_val_every_n_epoch=cfg.schedule.val_intervals,
accelerator='ddp',
log_every_n_steps=cfg.log.interval,
num_sanity_val_steps=0)
trainer.fit(task, train_dataloader, val_dataloader)
def main(config, model_path, output_path, input_shape=(320, 320)):
logger = Logger(-1, config.save_dir, False)
model = build_model(config.model)
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
load_model_weight(model, checkpoint, logger)
if config.model.arch.backbone.name == "RepVGG":
deploy_config = config.model
deploy_config.arch.backbone.update({"deploy": True})
deploy_model = build_model(deploy_config)
from nanodet.model.backbone.repvgg import repvgg_det_model_convert
model = repvgg_det_model_convert(model, deploy_model)
dummy_input = torch.autograd.Variable(
torch.randn(1, 3, input_shape[0], input_shape[1])
)
torch.onnx.export(
model,
dummy_input,
output_path,
verbose=True,
keep_initializers_as_inputs=True,
opset_version=11,
input_names=["data"],
output_names=["output"],
)
logger.log("finished exporting onnx ")
logger.log("start simplifying onnx ")
input_data = {"data": dummy_input.detach().cpu().numpy()}
model_sim, flag = onnxsim.simplify(output_path, input_data=input_data)
if flag:
onnx.save(model_sim, output_path)
logger.log("simplify onnx successfully")
else:
logger.log("simplify onnx failed")
def main(args):
load_config(cfg, args.config)
local_rank = -1
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
cfg.defrost()
timestr = datetime.datetime.now().__format__('%Y%m%d%H%M%S')
cfg.save_dir = os.path.join(cfg.save_dir, timestr)
mkdir(local_rank, cfg.save_dir)
logger = Logger(local_rank, cfg.save_dir)
assert args.task in ['val', 'test']
cfg.update({'test_mode': args.task})
logger.log('Setting up data...')
val_dataset = build_dataset(cfg.data.val, args.task)
val_dataloader = torch.utils.data.DataLoader(
val_dataset,
batch_size=cfg.device.batchsize_per_gpu,
shuffle=False,
num_workers=cfg.device.workers_per_gpu,
pin_memory=True,
collate_fn=collate_function,
drop_last=True)
evaluator = build_evaluator(cfg, val_dataset)
logger.log('Creating model...')
task = TrainingTask(cfg, evaluator)
ckpt = torch.load(args.model)
if 'pytorch-lightning_version' not in ckpt:
warnings.warn(
'Warning! Old .pth checkpoint is deprecated. '
'Convert the checkpoint with tools/convert_old_checkpoint.py ')
ckpt = convert_old_model(ckpt)
task.load_state_dict(ckpt['state_dict'])
trainer = pl.Trainer(
default_root_dir=cfg.save_dir,
gpus=cfg.device.gpu_ids,
accelerator='ddp',
log_every_n_steps=cfg.log.interval,
num_sanity_val_steps=0,
)
logger.log('Starting testing...')
trainer.test(task, val_dataloader)
def __init__(self, model_path, cfg_path, *args, **kwargs):
from nanodet.model.arch import build_model
from nanodet.util import Logger, cfg, load_config, load_model_weight
super(NanoDetTorch, self).__init__(*args, **kwargs)
print(f'Using PyTorch as inference backend')
print(f'Using weight: {model_path}')
# load model
self.model_path = model_path
self.cfg_path = cfg_path
load_config(cfg, cfg_path)
self.logger = Logger(-1, cfg.save_dir, False)
self.model = build_model(cfg.model)
checkpoint = self.torch.load(model_path, map_location=lambda storage, loc: storage)
load_model_weight(self.model, checkpoint, self.logger)
def main(config,
model_path,
output_path,
input_shape=(320, 320),
batch_size=1):
logger = Logger(-1, config.save_dir, False)
model = build_model(config.model)
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
load_model_weight(model, checkpoint, logger)
dummy_input = torch.autograd.Variable(
torch.randn(batch_size, 3, input_shape[0], input_shape[1]))
torch.onnx.export(model,
dummy_input,
output_path,
verbose=True,
keep_initializers_as_inputs=True,
opset_version=12)
onnx_model = onnx.load(output_path) # load onnx model
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, output_path)
print('finished exporting onnx ')
with torch.no_grad():
results = self.model.inference(meta)
return meta, results
def visualize(self, dets, meta, class_names, score_thres, wait=0):
time1 = time.time()
self.model.head.show_result(meta['raw_img'],
dets,
class_names,
score_thres=score_thres,
show=True)
print('viz time: {:.3f}s'.format(time.time() - time1))
if __name__ == '__main__':
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
load_config(cfg, config_path)
logger = Logger(-1, use_tensorboard=False)
predictor = Predictor(cfg, model_path, logger, device='cuda:0')
logger.log('Press "Esc" to exit.')
cap = cv2.VideoCapture(cam_id) #str for video
while True:
ret_val, frame = cap.read()
if ret_val == False:
continue #skip if capture fail
meta, res = predictor.inference(frame)
predictor.visualize(res, meta, cfg.class_names, 0.35)
ch = cv2.waitKey(1)
if ch == 27:
break
def main():
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind(('127.0.0.1', 8000))
server.listen(5)
print("waiting msg ...")
conn, clint_add = server.accept()
args = parse_args()
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
load_config(cfg, args.config)
logger = Logger(-1, use_tensorboard=False)
predictor = Predictor(cfg, args.model, logger, device='cuda:0')
logger.log('Press "Esc", "q" or "Q" to exit.')
if args.demo == 'image':
if os.path.isdir(args.path):
files = get_image_list(args.path)
else:
files = [args.path]
files.sort()
for image_name in files:
meta, res = predictor.inference(image_name)
predictor.visualize(res, meta, cfg.class_names, 0.35)
ch = cv2.waitKey(0)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
elif args.demo == 'video' or args.demo == 'webcam':
pipeline = rs.pipeline()
# 创建 config 对象：
config = rs.config()
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
# Start streaming
pipeline.start(config)
#cap = cv2.VideoCapture(args.path if args.demo == 'video' else args.camid)
while True:
next_frames = pipeline.wait_for_frames()
get_next_color_frame = next_frames.get_color_frame()
frame = np.asanyarray(get_next_color_frame.get_data())
meta, res = predictor.inference(frame)
predictor.visualize(res, meta, cfg.class_names, 0.70)
all_box = []
for label in res:
for bbox in res[label]:
score = bbox[-1]
if score > 0.70:
x0, y0, x1, y1 = [int(i) for i in bbox[:4]]
all_box.append([label, x0, y0, x1, y1, score])
all_box.sort(key=lambda v: v[5])
send_data_byte = bytes(0)
time.sleep(0.005)
if len(all_box) == 0:
leftup_rightdown_corner = [-1, 0, 0, 0, 0, time.time(), 'b']
for i in range(len(leftup_rightdown_corner)):
#print(pickup_leftup_rightdown_corner[i])
pickup_senddata = str(leftup_rightdown_corner[i]) + ','
# print(pickup_senddata.encode())
send_data_byte += pickup_senddata.encode()
# print(send_data_byte)
conn.send(send_data_byte)
else:
zzw = all_box[-1]
label, x0, y0, x1, y1, score = zzw
leftup_rightdown_corner = [1, x0, y0, x1, y1, time.time(), 'a']
for i in range(len(leftup_rightdown_corner)):
target_senddata = str(leftup_rightdown_corner[i]) + ','
send_data_byte += target_senddata.encode()
conn.send(send_data_byte)
ch = cv2.waitKey(1)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
def startNanodetTrain(self):
#加载配置文件
load_config(cfg, self.nanoTrainConfig['cfg'])
#判断分布式训练当中该主机的角色
local_rank = int(self.nanoTrainConfig["local_rank"])
# torch.backends.cudnn.enabled = True
# torch.backends.cudnn.benchmark = True
mkdir(local_rank, self.nanoTrainConfig["save_dir"])
logger = Logger(local_rank, self.nanoTrainConfig["save_dir"])
if self.nanoTrainConfig.keys().__contains__("seed"):
logger.log('Set random seed to {}'.format(
self.nanoTrainConfig['seed']))
self.init_seeds(self.nanoTrainConfig['seed'])
#1.创建模型
model = build_model(cfg.model)
model = model.cpu()
#2.加载数据
logger.log('Setting up data...')
train_dataset = build_dataset(cfg.data.train, 'train',
self.nanoTrainConfig)
val_dataset = build_dataset(cfg.data.val, 'test', self.nanoTrainConfig)
if len(cfg.device.gpu_ids) > 1:
print('rank = ', local_rank)
num_gpus = torch.cuda.device_count()
torch.cuda.set_device(local_rank % num_gpus)
dist.init_process_group(backend='nccl')
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.device.batchsize_per_gpu,
num_workers=cfg.device.workers_per_gpu,
pin_memory=True,
collate_fn=collate_function,
sampler=train_sampler,
drop_last=True)
else:
print("加载数据...")
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.device.batchsize_per_gpu,
shuffle=True,
num_workers=cfg.device.workers_per_gpu,
pin_memory=True,
collate_fn=collate_function,
drop_last=True)
val_dataloader = torch.utils.data.DataLoader(
val_dataset,
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True,
collate_fn=collate_function,
drop_last=True)
trainer = build_trainer(local_rank, cfg, model, logger)
if 'load_model' in cfg.schedule:
trainer.load_model(cfg)
if 'resume' in cfg.schedule:
trainer.resume(cfg)
evaluator = build_evaluator(cfg, val_dataset)
logger.log('Starting training...')
trainer.run(train_dataloader, val_dataloader, evaluator,
self.nanoTrainConfig)
def run(args):
"""
:param args:
:return:
"""
load_config(cfg, args.config)
local_rank = int(args.local_rank) # what's this?
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
mkdir(local_rank, cfg.save_dir)
logger = Logger(local_rank, cfg.save_dir)
if args.seed is not None:
logger.log('Set random seed to {}'.format(args.seed))
init_seeds(args.seed)
logger.log('Creating model...')
model = build_model(cfg.model)
logger.log('Setting up data...')
train_dataset = build_dataset(cfg.data.train, 'train') # build_dataset(cfg.data.train, 'train')
val_dataset = build_dataset(cfg.data.val, 'test')
if len(cfg.device.gpu_ids) > 1: # More than one GPU(distributed training)
print('rank = ', local_rank)
num_gpus = torch.cuda.device_count()
torch.cuda.set_device(local_rank % num_gpus)
dist.init_process_group(backend='nccl')
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
if args.is_debug:
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=cfg.device.batchsize_per_gpu,
num_workers=0,
pin_memory=True,
collate_fn=collate_function,
sampler=train_sampler,
drop_last=True)
else:
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=cfg.device.batchsize_per_gpu,
num_workers=cfg.device.workers_per_gpu,
pin_memory=True,
collate_fn=collate_function,
sampler=train_sampler,
drop_last=True)
else:
if args.is_debug:
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=cfg.device.batchsize_per_gpu,
shuffle=True,
num_workers=0,
pin_memory=True,
collate_fn=collate_function,
drop_last=True)
else:
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=cfg.device.batchsize_per_gpu,
shuffle=True,
num_workers=cfg.device.workers_per_gpu,
pin_memory=True,
collate_fn=collate_function,
drop_last=True)
if args.is_debug:
val_data_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=0,
pin_memory=True,
collate_fn=collate_function, drop_last=True)
else:
val_data_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True,
collate_fn=collate_function, drop_last=True)
# -----
trainer = build_trainer(local_rank, cfg, model, logger)
if 'load_model' in cfg.schedule:
trainer.load_model(cfg)
if 'resume' in cfg.schedule:
trainer.resume(cfg)
# ----- Build a evaluator
evaluator = build_evaluator(cfg, val_dataset)
# evaluator = None
logger.log('Starting training...')
trainer.run(train_data_loader, val_data_loader, evaluator)
