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():
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):
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 = NanoDetLightningLogger(cfg.save_dir)
assert args.task in ["val", "test"]
cfg.update({"test_mode": args.task})
logger.info("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=naive_collate,
drop_last=False,
)
evaluator = build_evaluator(cfg.evaluator, val_dataset)
logger.info("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"])
if cfg.device.gpu_ids == -1:
logger.info("Using CPU training")
accelerator, devices = "cpu", None
else:
accelerator, devices = "gpu", cfg.device.gpu_ids
trainer = pl.Trainer(
default_root_dir=cfg.save_dir,
accelerator=accelerator,
devices=devices,
log_every_n_steps=cfg.log.interval,
num_sanity_val_steps=0,
logger=logger,
)
logger.info("Starting testing...")
trainer.test(task, val_dataloader)
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):
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(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 __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(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 test_logger():
tmp_dir = tempfile.TemporaryDirectory()
logger = NanoDetLightningLogger(tmp_dir.name)
writer = logger.experiment
assert isinstance(writer, SummaryWriter)
logger.info("test")
logger.log_hyperparams({"lr": 1})
logger.log_metrics({"mAP": 30.1}, 1)
load_config(cfg, "./config/legacy_v0.x_configs/nanodet-m.yml")
logger.dump_cfg(cfg)
logger.finalize(None)
def getDownloadModelUrl(self, dmtvid):
# 查询该版本
modelVersion = detectModelTrainVersion.objects(
dmtvid=int(dmtvid),
state=ConstantUtils.DATA_STATUS_ACTIVE).first()
if modelVersion[
"inferencePlatform"] == ConstantUtils.MODEL_PLATFORM_SERVER:
return modelVersion.ckptModelSavePath
elif modelVersion[
"inferencePlatform"] == ConstantUtils.MODEL_PLATFORM_LITE:
# 获取该项目的weights路径
cfg_path = r"data/nanodet-self.yml"
model_path = str(modelVersion["ckptModelSavePath"])
load_config(cfg, cfg_path)
return liteConveter.convertToNCNN_Android_model(cfg, model_path)
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 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 test_config_files():
root_path = join(dirname(__file__), "../..")
cfg_folder = join(root_path, "config")
if not exists(cfg_folder):
raise FileNotFoundError("Cannot find config folder.")
cfg_paths = collect_files(cfg_folder, [".yml", ".yaml"])
for cfg_path in cfg_paths:
print(f"Start testing {cfg_path}")
config = copy.deepcopy(cfg)
# test load cfg
load_config(config, cfg_path)
assert "save_dir" in config
assert "model" in config
assert "data" in config
assert "device" in config
assert "schedule" in config
assert "log" in config
# test build model
model = build_model(config.model)
assert config.model.arch.name == model.__class__.__name__
print('finished exporting onnx ')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--cfg_file',
type=str,
default='./config/nanodet-m.yml',
help='config file path')
parser.add_argument('--weights_file',
type=str,
default='./weights/nanodet_m.pth',
help='weights file path')
parser.add_argument('--output_file',
type=str,
default='./nanodet-m.onnx',
help='onnx file path')
parser.add_argument('--img_size',
nargs='+',
type=int,
default=[320, 320],
help='image size')
parser.add_argument('--batch_size', type=int, default=1, help='batch size')
opt = parser.parse_args()
load_config(cfg, opt.cfg_file)
main(cfg,
opt.weights_file,
opt.output_file,
input_shape=opt.img_size,
batch_size=opt.batch_size)
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)
def test_lightning_training_task():
load_config(cfg, "./config/legacy_v0.x_configs/nanodet-m.yml")
task = TrainingTask(cfg)
runner = DummyRunner(task)
runner.test()
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 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 = NanoDetLightningLogger(cfg.save_dir)
logger.dump_cfg(cfg)
if args.seed is not None:
logger.info("Set random seed to {}".format(args.seed))
pl.seed_everything(args.seed)
logger.info("Setting up data...")
train_dataset = build_dataset(cfg.data.train, "train")
val_dataset = build_dataset(cfg.data.val, "test")
evaluator = build_evaluator(cfg.evaluator, 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=naive_collate,
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=naive_collate,
drop_last=False,
)
logger.info("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)
load_model_weight(task.model, ckpt, logger)
logger.info("Loaded model weight from {}".format(
cfg.schedule.load_model))
model_resume_path = (os.path.join(cfg.save_dir, "model_last.ckpt")
if "resume" in cfg.schedule else None)
accelerator = None if len(cfg.device.gpu_ids) <= 1 else "ddp"
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=accelerator,
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
logger=logger,
benchmark=True,
gradient_clip_val=cfg.get("grad_clip", 0.0),
)
trainer.fit(task, train_dataloader, val_dataloader)
def test_flops():
load_config(cfg, "./config/legacy_v0.x_configs/nanodet-m.yml")
model = build_model(cfg.model)
input_shape = (3, 320, 320)
get_model_complexity_info(model, input_shape)
help="Path to .ckpt model.")
parser.add_argument(
"--out_path",
type=str,
default="nanodet.torchscript.pth",
help="TorchScript model output path.",
)
parser.add_argument("--input_shape",
type=str,
default=None,
help="Model input shape.")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
cfg_path = args.cfg_path
model_path = args.model_path
out_path = args.out_path
input_shape = args.input_shape
load_config(cfg, cfg_path)
if input_shape is None:
input_shape = cfg.data.train.input_size
else:
input_shape = tuple(map(int, input_shape.split(",")))
assert len(input_shape) == 2
if model_path is None:
model_path = os.path.join(cfg.save_dir, "model_best/model_best.ckpt")
main(cfg, model_path, out_path, input_shape)
print("Model saved to:", out_path)
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)
