img_id='default')
# else:
# detection_result.append([0,0,0,0,0])
# cv2.imshow("detection",image_detection)
# cv2.waitKey(10)
return image_detection, detection_result
if __name__ == '__main__':
num_classes = 20
max_per_image = 100
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from models import get_pose_net
heads = {"hm": num_classes, "wh": 2, "reg": 2}
model = get_pose_net(50, heads, head_conv=64)
# model = load_model(model,"./models/model_origin.pth")
model = load_model(model, "./model_origin.pth")
model.cuda()
model.eval()
# video = cv2.VideoCapture("t640480_det_results.avi")
# # video = cv2.VideoCapture("MOT16-11.mp4")
# # Exit if video not opened.
# if not video.isOpened():
# print("Could not open video")
# sys.exit()
# # Read first frame.
# ok, frame = video.read()
def main(opt):
torch.manual_seed(opt.seed)
torch.backends.cudnn.benchmark = True
train_path = "./VOC/train.txt"
val_path = "./VOC/val.txt"
start_epoch = 0
imageshape = (384, 384)
val_loader = torch.utils.data.DataLoader(
listDataset(val_path, shape=imageshape, shuffle=False, train=False),
shuffle=False,
num_workers=0,
pin_memory=True,
batch_size=opt.batch_size,
)
train_loader = torch.utils.data.DataLoader(
listDataset(train_path, shape=imageshape, shuffle=True, train=True),
batch_size=opt.batch_size,
shuffle=True,
num_workers=0,
pin_memory=True,
)
print("the train_loader size is {}".format(len(train_loader)))
writer = SummaryWriter()
num_train_batches = int(len(train_loader))
num_val_batches = int(len(val_loader))
logger.info("the train batches is {}".format(num_train_batches))
logger.info("the val batches is {}".format(num_val_batches))
from models import get_pose_net
heads = {"hm": 20, "wh": 2, "reg": 2}
model = get_pose_net(18, heads, head_conv=64)
model.cuda()
criterion = CtdetLoss(opt)
# optimizer = torch.optim.SGD(model.parameters(),lr=opt.lr, momentum = 0.9,weight_decay=1e-4)
optimizer = torch.optim.Adam(model.parameters(), opt.lr, weight_decay=1e-3)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=opt.lr_step,
gamma=0.1)
config = {
'tensorboard': True,
'tensorboard_images': True,
'tensorboard_parameters': True,
}
best_val = 1000
# run test before start training
test(0, model, criterion, val_loader, config, writer)
for epoch in range(start_epoch, opt.num_epochs):
scheduler.step()
train(epoch, model, optimizer, criterion, train_loader, config, writer)
test_val = test(epoch, model, criterion, val_loader, config, writer)
state = OrderedDict([
# ('args', vars(args)),
('state_dict', model.state_dict()),
('optimizer', optimizer.state_dict()),
('epoch', epoch),
('angle_error', test_val),
])
if test_val < best_val:
outdir = "./"
model_path = os.path.join(outdir, opt.model_path)
torch.save(state, model_path)
torch.save(state, "./last.pth")