def debug(self, batch, output, iter_id):
opt = self.opt
reg = output['reg'] if opt.reg_offset else None
dets = ctdet_decode(output['hm'],
output['wh'],
reg=reg,
cat_spec_wh=opt.cat_spec_wh,
K=opt.K)
dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
dets[:, :, :4] *= opt.down_ratio
dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])
dets_gt[:, :, :4] *= opt.down_ratio
for i in range(1):
debugger = Debugger(dataset=opt.dataset,
ipynb=(opt.debug == 3),
theme=opt.debugger_theme)
img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = np.clip(((img * opt.std + opt.mean) * 255.), 0,
255).astype(np.uint8)
pred = debugger.gen_colormap(
output['hm'][i].detach().cpu().numpy())
gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'pred_hm')
debugger.add_blend_img(img, gt, 'gt_hm')
debugger.add_img(img, img_id='out_pred')
for k in range(len(dets[i])):
if dets[i, k, 4] > opt.center_thresh:
debugger.add_coco_bbox(dets[i, k, :4],
dets[i, k, -1],
dets[i, k, 4],
img_id='out_pred')
debugger.add_img(img, img_id='out_gt')
for k in range(len(dets_gt[i])):
if dets_gt[i, k, 4] > opt.center_thresh:
debugger.add_coco_bbox(dets_gt[i, k, :4],
dets_gt[i, k, -1],
dets_gt[i, k, 4],
img_id='out_gt')
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir,
prefix='{}'.format(iter_id))
else:
debugger.show_all_imgs(pause=False)
def debug(self, batch, output, iter_id):
opt = self.opt
reg = output['reg'] if opt.reg_offset else None
hm_hp = output['hm_hp'] if opt.hm_hp else None
hp_offset = output['hp_offset'] if opt.reg_hp_offset else None
dets = multi_pose_decode(
output['hm'], output['wh'], output['hps'],
reg=reg, hm_hp=hm_hp, hp_offset=hp_offset, K=opt.K)
dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
dets[:, :, :4] *= opt.input_res / opt.output_res
dets[:, :, 5:5 + Num_Key_Point * 2] *= opt.input_res / opt.output_res
dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])
dets_gt[:, :, :4] *= opt.input_res / opt.output_res
dets_gt[:, :, 5:5 + Num_Key_Point * 2] *= opt.input_res / opt.output_res
for i in range(1):
debugger = Debugger(
dataset=opt.dataset, ipynb=(opt.debug == 3), theme=opt.debugger_theme)
img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = np.clip(((img * opt.std + opt.mean) * 255.), 0, 255).astype(np.uint8)
pred = debugger.gen_colormap(output['hm'][i].detach().cpu().numpy())
gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'pred_hm')
debugger.add_blend_img(img, gt, 'gt_hm')
debugger.add_img(img, img_id='out_pred')
for k in range(len(dets[i])):
if dets[i, k, 4] > opt.center_thresh:
debugger.add_coco_bbox(dets[i, k, :4], dets[i, k, -1],
dets[i, k, 4], img_id='out_pred')
debugger.add_coco_hp(dets[i, k, 5:5 + Num_Key_Point * 2], img_id='out_pred')
debugger.add_img(img, img_id='out_gt')
for k in range(len(dets_gt[i])):
if dets_gt[i, k, 4] > opt.center_thresh:
debugger.add_coco_bbox(dets_gt[i, k, :4], dets_gt[i, k, -1],
dets_gt[i, k, 4], img_id='out_gt')
debugger.add_coco_hp(dets_gt[i, k, 5:5 + Num_Key_Point * 2], img_id='out_gt')
if opt.hm_hp:
pred = debugger.gen_colormap_hp(output['hm_hp'][i].detach().cpu().numpy())
gt = debugger.gen_colormap_hp(batch['hm_hp'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'pred_hmhp')
debugger.add_blend_img(img, gt, 'gt_hmhp')
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir, prefix='{}'.format(iter_id))
else:
debugger.show_all_imgs(pause=True)
def _debug(image, t_heat, l_heat, b_heat, r_heat, ct_heat):
debugger = Debugger(num_classes=1)
k = 0
t_heat = torch.sigmoid(t_heat)
l_heat = torch.sigmoid(l_heat)
b_heat = torch.sigmoid(b_heat)
r_heat = torch.sigmoid(r_heat)
aggr_weight = 0.1
t_heat = _h_aggregate(t_heat, aggr_weight=aggr_weight)
l_heat = _v_aggregate(l_heat, aggr_weight=aggr_weight)
b_heat = _h_aggregate(b_heat, aggr_weight=aggr_weight)
r_heat = _v_aggregate(r_heat, aggr_weight=aggr_weight)
t_heat[t_heat > 1] = 1
l_heat[l_heat > 1] = 1
b_heat[b_heat > 1] = 1
r_heat[r_heat > 1] = 1
ct_heat = torch.sigmoid(ct_heat)
t_hm = debugger.gen_colormap(t_heat[k].cpu().data.numpy())
l_hm = debugger.gen_colormap(l_heat[k].cpu().data.numpy())
b_hm = debugger.gen_colormap(b_heat[k].cpu().data.numpy())
r_hm = debugger.gen_colormap(r_heat[k].cpu().data.numpy())
ct_hm = debugger.gen_colormap(ct_heat[k].cpu().data.numpy())
hms = np.maximum(np.maximum(t_hm, l_hm), np.maximum(b_hm, r_hm))
if image is not None:
mean = np.array([0.40789654, 0.44719302, 0.47026115],
dtype=np.float32).reshape(3, 1, 1)
std = np.array([0.28863828, 0.27408164, 0.27809835],
dtype=np.float32).reshape(3, 1, 1)
img = (image[k].cpu().data.numpy() * std + mean) * 255
img = img.astype(np.uint8).transpose(1, 2, 0)
debugger.add_img(img, 'img')
debugger.add_blend_img(img, t_hm, 't_hm')
debugger.add_blend_img(img, l_hm, 'l_hm')
debugger.add_blend_img(img, b_hm, 'b_hm')
debugger.add_blend_img(img, r_hm, 'r_hm')
debugger.add_blend_img(img, hms, 'extreme')
debugger.add_blend_img(img, ct_hm, 'center')
debugger.show_all_imgs(pause=False)
def main():
opt = opts().parse()
#if opt.loadModel != 'none':
model = AlexNet(ref.nJoints).cuda()
model.load_state_dict(torch.load("save.model"))
for (i, filename) in enumerate(os.listdir("./testimages/")):
img = cv2.imread("./testimages/" + filename)
c = np.ones(2) * ref.h36mImgSize / 2
s = ref.h36mImgSize * 1.0
img2 = Crop(img, c, s, 0, ref.inputRes) / 256.
input = torch.from_numpy(img2)
input = input.contiguous().view(1, input.size(0), input.size(1),
input.size(2))
print(input.size())
input_var = torch.autograd.Variable(input).float().cuda()
output = model(input_var)
print(output.size())
reg = (output.data).cpu().numpy() #.reshape(pred.shape[0], 1)
four = lambda t: t * 4.57
fourfunc = np.vectorize(four)
reg = fourfunc(reg)
print(reg)
debugger = Debugger()
debugger.addImg(
(input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8))
debugger.addPoint2D(reg, (255, 0, 0))
#debugger.addPoint3D(np.concatenate([pred, (reg + 1) / 2. * 256], axis = 1))
debugger.saveImg(path="./result/" + filename)
np.set_printoptions(threshold=np.inf, linewidth=np.inf)
with open("./result/" + filename[:-4] + ".out", 'w') as f:
f.write(np.array2string(reg, separator=', '))
"""
def debug(self, batch, output, iter_id):
opt = self.opt
wh = output['wh'] if opt.reg_bbox else None
reg = output['reg'] if opt.reg_offset else None
dets = car_pose_6dof_decode(output['hm'],
output['rot'],
output['dep'],
wh=wh,
reg=reg,
K=opt.K)
# x, y, score, r1-r4, depth, wh?, cls
dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
c = batch['meta']['c'].detach().numpy()
s = batch['meta']['s'].detach().numpy()
calib = batch['meta']['calib'].detach().numpy()
# yaw, pitch, roll, x, y, z, wh?, score
dets_pred = car_6dof_post_process(dets.copy(), c, s, calib, opt)
dets_gt = car_6dof_post_process(
batch['meta']['gt_det'].detach().numpy().copy(), c, s, calib, opt)
car_name = car_models.models[0].name
car_model = self.models[car_name]
for i in range(1):
debugger = Debugger(dataset=opt.dataset,
ipynb=(opt.debug == 3),
theme=opt.debugger_theme)
img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = ((img * self.opt.std + self.opt.mean) * 255.).astype(
np.uint8)
if opt.debug_heatmap:
pred = debugger.gen_colormap(
output['hm'][i].detach().cpu().numpy())
gt = debugger.gen_colormap(
batch['hm'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'hm_pred')
debugger.add_blend_img(img, gt, 'hm_gt')
if opt.render_cars:
debugger.render_cars(img, dets_pred[0], car_model, c, s, calib,
opt, '3d_pred')
debugger.render_cars(img, dets_gt[0], car_model, c, s, calib,
opt, '3d_gt')
if opt.debug == 4:
prefix = '{}_{}_'.format(iter_id, batch['meta']['img_id'][0])
debugger.save_all_imgs(opt.debug_dir, prefix=prefix)
else:
debugger.show_all_imgs(pause=True)
def debug(self, batch, output, iter_id):
opt = self.opt
detections = self.decode(output['hm_t'], output['hm_l'],
output['hm_b'], output['hm_r'],
output['hm_c']).detach().cpu().numpy()
detections[:, :, :4] *= opt.input_res / opt.output_res
for i in range(1):
dataset = opt.dataset
if opt.dataset == 'yolo':
dataset = opt.names
debugger = Debugger(dataset=dataset,
ipynb=(opt.debug == 3),
theme=opt.debugger_theme)
pred_hm = np.zeros((opt.input_res, opt.input_res, 3),
dtype=np.uint8)
gt_hm = np.zeros((opt.input_res, opt.input_res, 3), dtype=np.uint8)
img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = ((img * self.opt.std + self.opt.mean) * 255.).astype(
np.uint8)
for p in self.parts:
tag = 'hm_{}'.format(p)
pred = debugger.gen_colormap(
output[tag][i].detach().cpu().numpy())
gt = debugger.gen_colormap(
batch[tag][i].detach().cpu().numpy())
if p != 'c':
pred_hm = np.maximum(pred_hm, pred)
gt_hm = np.maximum(gt_hm, gt)
if p == 'c' or opt.debug > 2:
debugger.add_blend_img(img, pred, 'pred_{}'.format(p))
debugger.add_blend_img(img, gt, 'gt_{}'.format(p))
debugger.add_blend_img(img, pred_hm, 'pred')
debugger.add_blend_img(img, gt_hm, 'gt')
debugger.add_img(img, img_id='out')
for k in range(len(detections[i])):
if detections[i, k, 4] > 0.1:
debugger.add_coco_bbox(detections[i, k, :4],
detections[i, k, -1],
detections[i, k, 4],
img_id='out')
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir,
prefix='{}'.format(iter_id))
else:
debugger.show_all_imgs(pause=True)
def show_results(self, image, gts, dets, save_dir, img_name):
debugger = Debugger(dataset='dota',
ipynb=(self.opt.debug == 3),
theme='white')
debugger.add_img(image, img_name)
for j in dets:
for bbox in dets[j]:
if bbox[5] > 0.01:
debugger.add_rbbox(bbox[:5],
j - 1,
bbox[5],
img_id=img_name)
for ann in gts:
bbox = ann['rbbox']
cat_id = ann['category_id']
debugger.add_rbbox(bbox, cat_id - 1, 1, img_id=img_name, gt=True)
save_dir = os.path.join(save_dir, 'voc_results')
debugger.save_all_imgs(save_dir)
def demo(opt):
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
# even in default o mode, debug = 1
# opt.debug = max(opt.debug, 1) # 1: only show the final detection results
# update opt info with dataset
Dataset = dataset_factory[opt.dataset]
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
pprint(vars(opt))
Detector = detector_factory[opt.task] # choose a detector
detector = Detector(opt)
debugger = Debugger(dataset=opt.dataset,
ipynb=False,
theme=opt.debugger_theme) # debugger.imgs is a dict
image_names, video_names = [], []
if os.path.isdir(opt.demo): # image dir or video drr
ls = os.listdir(opt.demo)
for file_name in sorted(ls):
ext = file_name[file_name.rfind('.') + 1:].lower() # img ext
if ext in image_ext:
image_names.append(os.path.join(opt.demo, file_name))
if ext in video_ext:
video_names.append(os.path.join(opt.demo, file_name))
elif opt.demo[opt.demo.rfind('.') + 1:].lower() in video_ext: # one video
video_names = [opt.demo]
elif opt.demo[opt.demo.rfind('.') + 1:].lower() in image_ext: # one image
image_names = [opt.demo]
elif opt.demo == 'webcam': # cam
pass
for image_name in image_names:
print(image_name)
process_img(image_name, detector)
opt.debug = 0 # when test video, use draw_res_box
for video_name in video_names:
print(video_name)
process_video(video_name, detector, debugger, Dataset.num_classes)
def main():
opt = opts().parse()
if opt.loadModel != 'none':
model = torch.load(opt.loadModel).cuda()
else:
model = torch.load('hgreg-3d.pth').cuda()
img = cv2.imread(opt.demo)
input = torch.from_numpy(img.transpose(2, 0, 1)).float() / 256.
input = input.view(1, input.size(0), input.size(1), input.size(2))
input_var = torch.autograd.Variable(input).float().cuda()
output = model(input_var)
pred = getPreds((output[-2].data).cpu().numpy())[0] * 4
reg = (output[-1].data).cpu().numpy().reshape(pred.shape[0], 1)
debugger = Debugger()
debugger.addImg((input[0].numpy().transpose(1, 2, 0)*256).astype(np.uint8))
debugger.addPoint2D(pred, (255, 0, 0))
debugger.addPoint3D(np.concatenate([pred, (reg + 1) / 2. * 256], axis = 1))
debugger.showImg(pause = True)
debugger.show3D()
def show_results(self, save_dir):
with open(os.path.join(save_dir, "results.json")) as f:
data = json.load(f)
debugger = Debugger(dataset=self.opt.dataset,
ipynb=(self.opt.debug == 3),
theme=self.opt.debugger_theme,
class_names=self.opt.class_names)
for i, img_details in enumerate(data):
if (data[i]['score'] > self.opt.vis_thresh):
img_path = os.path.join("../data/coco/test2019",
str(img_details["image_id"]) +
".jpg") # TODO
image = cv2.imread(img_path)
debugger.add_img(image, img_id='ctdet')
bbox = data[i]['bbox']
debugger.add_coco_bbox(bbox,
data[i]['category_id'] - 1,
data[i]['score'],
img_id='ctdet')
debugger.show_all_imgs(pause=True)
def check_logic():
preproc = nn.ModuleList([model.conv1_, model.bn1, model.relu, model.r1, model.maxpool, model.r4, model.r5 ])
hg = model.hourglass[0]
lower_hg = getEncoder(hg)
data_loader = torch.utils.data.DataLoader(
H36M(opts, 'val'),
batch_size = 1,
shuffle = False,
num_workers = int(ref.nThreads)
)
for k, (input, target) in enumerate(data_loader):
if(k>nSamples):
break
input_var = torch.autograd.Variable(input).float().cuda()
for mod in preproc:
input_var = mod(input_var)
for mod in lower_hg:
input_var = mod(input_var)
#decode
ups = input_var
upper_hg = nn.ModuleList(getDecoder(hg))
for mod in upper_hg:
ups = mod(ups)
Residual = model.Residual
for j in range(nModules):
ups = Residual[j](ups)
lin_ = model.lin_
ups = lin_[0](ups)
tmpOut = model.tmpOut
ups = tmpOut[0](ups)
pred = eval.getPreds(ups.data.cpu().numpy())*4
gt = eval.getPreds(target.cpu().numpy()) * 4
# init = getPreds(input.numpy()[:, 3:])
debugger = Debugger()
img = (input[0].numpy()[:3].transpose(1, 2, 0)*256).astype(np.uint8).copy()
print(img.shape)
debugger.addImg(img)
debugger.addPoint2D(pred[0], (255, 0, 0))
debugger.addPoint2D(gt[0], (0, 0, 255))
# debugger.addPoint2D(init[0], (0, 255, 0))
debugger.showAllImg(pause = True)
def demo_image(image, model, opt):
s = max(image.shape[0], image.shape[1]) * 1.0
c = np.array([image.shape[1] / 2., image.shape[0] / 2.], dtype=np.float32)
trans_input = get_affine_transform(
c, s, 0, [opt.input_w, opt.input_h])
inp = cv2.warpAffine(image, trans_input, (opt.input_w, opt.input_h),
flags=cv2.INTER_LINEAR)
inp = (inp / 255. - mean) / std
inp = inp.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32)
inp = torch.from_numpy(inp).to(opt.device)
out = model(inp)[-1]
pred = get_preds(out['hm'].detach().cpu().numpy())[0]
pred = transform_preds(pred, c, s, (opt.output_w, opt.output_h))
pred_3d = get_preds_3d(out['hm'].detach().cpu().numpy(),
out['depth'].detach().cpu().numpy())[0]
debugger = Debugger()
debugger.add_img(image)
debugger.add_point_2d(pred, (255, 0, 0))
debugger.add_point_3d(pred_3d, 'b')
debugger.show_all_imgs(pause=False)
debugger.show_3d()
def __init__(self, config):
self.weight_path = config['predictor']['weight']
self.scale = config['predictor']['scale']
self.threshold = config['predictor']['threshold']
self.max_obj_predict = config['dataset']['max_object']
self.num_classes = config['dataset']['num_classes']
self.arch = config['model']['arch']
self.heads = config['model']['heads']
self.head_conv = config['model']['head_conv']
self.model = create_model(arch=self.arch,
heads=self.heads,
head_conv=self.head_conv)
self.model = load_model(self.model, self.weight_path)
if torch.cuda.is_available():
self.model.cuda()
self.model.eval()
self.dt = BaseDetector(config)
self.debugger = Debugger(num_classes=self.num_classes)
def debug_for_polygon(self, batch, output, iter_id):
opt = self.opt
output = output[0]
batch = batch[0]
for i in range(1):
debugger = Debugger(dataset=opt.dataset,
ipynb=(opt.debug == 3),
theme=opt.debugger_theme)
img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = np.clip(((img * opt.std + opt.mean) * 255.), 0,
255).astype(np.uint8)
pred = debugger.gen_colormap(
output['hm'][i].detach().cpu().numpy())
gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'pred_hm')
debugger.add_blend_img(img, gt, 'gt_hm')
debugger.add_img(img, img_id='out_pred')
debugger.add_img(img, img_id='out_gt')
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir,
prefix='{}'.format(iter_id))
else:
debugger.show_all_imgs(pause=True)
def generate(imageName):
process_image(imageName)
model = torch.load('../model/Stage3/model_10.pth',
map_location=lambda storage, loc: storage)
img = cv2.imread(imageName)
input = torch.from_numpy(img.transpose(2, 0, 1)).float() / 256.
input = input.view(1, input.size(0), input.size(1), input.size(2))
input_var = torch.autograd.Variable(input).float()
output = model(input_var)
print(output[-2].data[0][-2].shape)
pred = getPreds((output[-2].data).cpu().numpy())[0] * 4
reg = (output[-1].data).cpu().numpy().reshape(pred.shape[0], 1)
print(pred, (reg + 1) / 2. * 256)
debugger = Debugger()
debugger.addImg(
(input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8))
debugger.addPoint2D(pred, (255, 0, 0))
debugger.addPoint3D(np.concatenate([pred, (reg + 1) / 2. * 256], axis=1))
debugger.showImg(pause=True)
debugger.show3D()
'''
def main():
pickle.load = partial(pickle.load, encoding="latin1")
pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1")
opt = opts().parse()
if opt.loadModel != 'none':
model = torch.load(opt.loadModel).cuda()
else:
model = torch.load('../../tr_models/hgreg-3d.pth').cuda()
#opt.demo has the path to dir containing frames of demo video
all_frames = os.listdir(opt.demo)
n_frames = len(all_frames)
#specifics
dir_name = opt.demo.split('/')[-1]
save_path = '../../output/demo/'+dir_name
try:
os.makedirs(save_path)
except OSError:
pass
for idx, frame in enumerate(all_frames):
print('processing frame {}'.format(idx))
img = cv2.imread(opt.demo+'/'+frame)
input = torch.from_numpy(img.transpose(2, 0, 1)).float() / 256.
input = input.view(1, input.size(0), input.size(1), input.size(2))
input_var = torch.autograd.Variable(input).float().cuda()
output = model(input_var)
pred = getPreds((output[-2].data).cpu().numpy())[0] * 4
reg = (output[-1].data).cpu().numpy().reshape(pred.shape[0], 1)
debugger = Debugger()
debugger.addImg((input[0].numpy().transpose(1, 2, 0)*256).astype(np.uint8))
debugger.addPoint2D(pred, (255, 0, 0))
debugger.addPoint3D(np.concatenate([pred, (reg + 1) / 2. * 256], axis = 1))
# debugger.showImg(pause = True)
debugger.saveImg(path=save_path+'/frame{}.jpg'.format(idx))
debugger.save3D(path=save_path+'/frame_p3d{}.jpg'.format(idx))
print('frame {} done'.format(idx))
def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device("cuda")
else:
opt.device = torch.device("cpu")
print("Creating model...")
self.model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
self.model = load_model(self.model, opt.load_model, opt)
self.model = self.model.to(opt.device)
self.model.eval()
self.opt = opt
self.trained_dataset = get_dataset(opt.dataset)
self.mean = np.array(self.trained_dataset.mean,
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array(self.trained_dataset.std,
dtype=np.float32).reshape(1, 1, 3)
# self.pause = not opt.no_pause
self.rest_focal_length = (self.trained_dataset.rest_focal_length
if self.opt.test_focal_length < 0 else
self.opt.test_focal_length)
self.flip_idx = self.trained_dataset.flip_idx
self.cnt = 0
self.pre_images = None
self.pre_image_ori = None
self.dataset = opt.dataset
if self.dataset == "nuscenes":
self.tracker = {}
for class_name in NUSCENES_TRACKING_NAMES:
self.tracker[class_name] = Tracker(opt, self.model)
else:
self.tracker = Tracker(opt, self.model)
self.debugger = Debugger(opt=opt, dataset=self.trained_dataset)
self.img_height = 100
self.img_width = 100
def run(self, image_or_path_or_tensor, meta=None):
load_time, pre_time, net_time, dec_time, post_time = 0, 0, 0, 0, 0
merge_time, tot_time = 0, 0
debugger = Debugger((self.cfg.DEBUG==3), theme=self.cfg.DEBUG_THEME,
num_classes=self.cfg.MODEL.NUM_CLASSES, dataset=self.cfg.SAMPLE_METHOD, down_ratio=self.cfg.MODEL.DOWN_RATIO)
start_time = time.time()
pre_processed = False
if isinstance(image_or_path_or_tensor, np.ndarray):
image = image_or_path_or_tensor
elif type(image_or_path_or_tensor) == type (''):
image = cv2.imread(image_or_path_or_tensor)
else:
image = image_or_path_or_tensor['image'][0].numpy()
pre_processed_images = image_or_path_or_tensor
pre_processed = True
loaded_time = time.time()
load_time += (loaded_time - start_time)
detections = []
for scale in self.scales:
scale_start_time = time.time()
if not pre_processed:
images, meta = self.pre_process(image, scale, meta)
else:
images = pre_processed_images['images'][scale][0]
meta = pre_processed_images['meta'][scale]
meta = {k: v.numpy()[0] for k, v in meta.items()}
images = images.to(torch.device('cuda'))
torch.cuda.synchronize()
pre_process_time = time.time()
pre_time += pre_process_time - scale_start_time
output, dets, forward_time = self.process(images, return_time=True)
torch.cuda.synchronize()
net_time += forward_time - pre_process_time
decode_time = time.time()
dec_time += decode_time - forward_time
if self.cfg.DEBUG >= 2:
self.debug(debugger, images, dets, output, scale)
dets= self.post_process(dets, meta, scale)
torch.cuda.synchronize()
post_process_time = time.time()
post_time += post_process_time - decode_time
detections.append(dets)
results = self.merge_outputs(detections)
torch.cuda.synchronize()
end_time = time.time()
merge_time += end_time - post_process_time
tot_time += end_time - start_time
if self.cfg.DEBUG >= 1:
self.show_results(debugger, image, results)
return {'results': {1:results}, 'tot': tot_time, 'load': load_time,
'pre': pre_time, 'net': net_time, 'dec': dec_time,
'post': post_time, 'merge': merge_time}
def debug(self, batch, output, iter_id, dataset):
opt = self.opt
if "pre_hm" in batch:
output.update({"pre_hm": batch["pre_hm"]})
dets = generic_decode(output, K=opt.K, opt=opt)
for k in dets:
dets[k] = dets[k].detach().cpu().numpy()
dets_gt = batch["meta"]["gt_det"]
for i in range(1):
debugger = Debugger(opt=opt, dataset=dataset)
img = batch["image"][i].detach().cpu().numpy().transpose(1, 2, 0)
img = np.clip(((img * dataset.std + dataset.mean) * 255.0), 0,
255).astype(np.uint8)
pred = debugger.gen_colormap(
output["hm"][i].detach().cpu().numpy())
gt = debugger.gen_colormap(batch["hm"][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, "pred_hm")
debugger.add_blend_img(img, gt, "gt_hm")
if "pre_img" in batch:
pre_img = batch["pre_img"][i].detach().cpu().numpy().transpose(
1, 2, 0)
pre_img = np.clip(
((pre_img * dataset.std + dataset.mean) * 255), 0,
255).astype(np.uint8)
debugger.add_img(pre_img, "pre_img_pred")
debugger.add_img(pre_img, "pre_img_gt")
if "pre_hm" in batch:
pre_hm = debugger.gen_colormap(
batch["pre_hm"][i].detach().cpu().numpy())
debugger.add_blend_img(pre_img, pre_hm, "pre_hm")
debugger.add_img(img, img_id="out_pred")
if "ltrb_amodal" in opt.heads:
debugger.add_img(img, img_id="out_pred_amodal")
debugger.add_img(img, img_id="out_gt_amodal")
# Predictions
for k in range(len(dets["scores"][i])):
if dets["scores"][i, k] > opt.vis_thresh:
debugger.add_coco_bbox(
dets["bboxes"][i, k] * opt.down_ratio,
dets["clses"][i, k],
dets["scores"][i, k],
img_id="out_pred",
)
if "ltrb_amodal" in opt.heads:
debugger.add_coco_bbox(
dets["bboxes_amodal"][i, k] * opt.down_ratio,
dets["clses"][i, k],
dets["scores"][i, k],
img_id="out_pred_amodal",
)
if "hps" in opt.heads and int(dets["clses"][i, k]) == 0:
debugger.add_coco_hp(dets["hps"][i, k] *
opt.down_ratio,
img_id="out_pred")
if "tracking" in opt.heads:
debugger.add_arrow(
dets["cts"][i][k] * opt.down_ratio,
dets["tracking"][i][k] * opt.down_ratio,
img_id="out_pred",
)
debugger.add_arrow(
dets["cts"][i][k] * opt.down_ratio,
dets["tracking"][i][k] * opt.down_ratio,
img_id="pre_img_pred",
)
# Ground truth
debugger.add_img(img, img_id="out_gt")
for k in range(len(dets_gt["scores"][i])):
if dets_gt["scores"][i][k] > opt.vis_thresh:
debugger.add_coco_bbox(
dets_gt["bboxes"][i][k] * opt.down_ratio,
dets_gt["clses"][i][k],
dets_gt["scores"][i][k],
img_id="out_gt",
)
if "ltrb_amodal" in opt.heads:
debugger.add_coco_bbox(
dets_gt["bboxes_amodal"][i, k] * opt.down_ratio,
dets_gt["clses"][i, k],
dets_gt["scores"][i, k],
img_id="out_gt_amodal",
)
if "hps" in opt.heads and (int(dets["clses"][i, k]) == 0):
debugger.add_coco_hp(dets_gt["hps"][i][k] *
opt.down_ratio,
img_id="out_gt")
if "tracking" in opt.heads:
debugger.add_arrow(
dets_gt["cts"][i][k] * opt.down_ratio,
dets_gt["tracking"][i][k] * opt.down_ratio,
img_id="out_gt",
)
debugger.add_arrow(
dets_gt["cts"][i][k] * opt.down_ratio,
dets_gt["tracking"][i][k] * opt.down_ratio,
img_id="pre_img_gt",
)
if "hm_hp" in opt.heads:
pred = debugger.gen_colormap_hp(
output["hm_hp"][i].detach().cpu().numpy())
gt = debugger.gen_colormap_hp(
batch["hm_hp"][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, "pred_hmhp")
debugger.add_blend_img(img, gt, "gt_hmhp")
if "rot" in opt.heads and "dim" in opt.heads and "dep" in opt.heads:
dets_gt = {k: dets_gt[k].cpu().numpy() for k in dets_gt}
calib = (batch["meta"]["calib"].detach().numpy()
if "calib" in batch["meta"] else None)
det_pred = generic_post_process(
opt,
dets,
batch["meta"]["c"].cpu().numpy(),
batch["meta"]["s"].cpu().numpy(),
output["hm"].shape[2],
output["hm"].shape[3],
self.opt.num_classes,
calib,
)
det_gt = generic_post_process(
opt,
dets_gt,
batch["meta"]["c"].cpu().numpy(),
batch["meta"]["s"].cpu().numpy(),
output["hm"].shape[2],
output["hm"].shape[3],
self.opt.num_classes,
calib,
)
debugger.add_3d_detection(
batch["meta"]["img_path"][i],
batch["meta"]["flipped"][i],
det_pred[i],
calib[i],
vis_thresh=opt.vis_thresh,
img_id="add_pred",
)
debugger.add_3d_detection(
batch["meta"]["img_path"][i],
batch["meta"]["flipped"][i],
det_gt[i],
calib[i],
vis_thresh=opt.vis_thresh,
img_id="add_gt",
)
debugger.add_bird_views(
det_pred[i],
det_gt[i],
vis_thresh=opt.vis_thresh,
img_id="bird_pred_gt",
)
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir,
prefix="{}".format(iter_id))
else:
debugger.show_all_imgs(pause=True)
def step(split, epoch, opt, data_loader, model, optimizer=None):
if split == 'train':
model.train()
else:
model.eval()
# crit = torch.nn.MSELoss()
# crit_3d = FusionLoss(opt.device, opt.weight_3d, opt.weight_var)
# crit_ocv = nn.BCEWithLogitsLoss()
crit_ocv = nn.CrossEntropyLoss()
# acc_idxs = data_loader.dataset.acc_idxs
# edges = data_loader.dataset.edges
# edges_3d = data_loader.dataset.edges_3d
# shuffle_ref = data_loader.dataset.shuffle_ref
# mean = data_loader.dataset.mean
# std = data_loader.dataset.std
# convert_eval_format = data_loader.dataset.convert_eval_format
# Loss, Loss3D = AverageMeter(), AverageMeter()
# Acc, MPJPE = AverageMeter(), AverageMeter()
Loss_ocv, Acc_ocv = AverageMeter(), AverageMeter()
data_time, batch_time = AverageMeter(), AverageMeter()
preds = []
time_str = ''
nIters = len(data_loader)
if opt.train_half:
nIters = nIters / 2
bar = Bar('{}'.format(opt.exp_id), max=nIters)
end = time.time()
for i, batch in enumerate(data_loader):
if i >= nIters:
break
data_time.update(time.time() - end)
# for k in batch:
# if k != 'meta':
# batch[k] = batch[k].cuda(device=opt.device, non_blocking=True)
# gt_2d = batch['meta']['pts_crop'].cuda(
# device=opt.device, non_blocking=True).float() / opt.output_h
img, ocv_gt, info = batch
if i == 0:
np.savez(split + '_debug.npz',
img=img.numpy(),
ocv_gt=ocv_gt.numpy(),
info=info)
img = img.cuda(device=opt.device, non_blocking=True)
ocv_gt = ocv_gt.cuda(device=opt.device, non_blocking=True)
output = model(img)
# loss = crit(output[-1]['hm'], batch['target'])
# loss_3d = crit_3d(
# output[-1]['depth'], batch['reg_mask'], batch['reg_ind'],
# batch['reg_target'],gt_2d)
# for k in range(opt.num_stacks - 1):
# loss += crit(output[k], batch['target'])
# loss_3d = crit_3d(
# output[-1]['depth'], batch['reg_mask'], batch['reg_ind'],
# batch['reg_target'], gt_2d)
# loss += loss_3d
# loss = crit_ocv(output, ocv_gt)
loss = crit_ocv(output, torch.argmax(ocv_gt, 1))
preds = torch.argmax(output, 1)
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
# else:
# input_ = batch['input'].cpu().numpy().copy()
# input_[0] = flip(input_[0]).copy()[np.newaxis, ...]
# input_flip_var = torch.from_numpy(input_).cuda(
# device=opt.device, non_blocking=True)
# output_flip_ = model(input_flip_var)
# output_flip = shuffle_lr(
# flip(output_flip_[-1]['hm'].detach().cpu().numpy()[0]), shuffle_ref)
# output_flip = output_flip.reshape(
# 1, opt.num_output, opt.output_h, opt.output_w)
# output_depth_flip = shuffle_lr(
# flip(output_flip_[-1]['depth'].detach().cpu().numpy()[0]), shuffle_ref)
# output_depth_flip = output_depth_flip.reshape(
# 1, opt.num_output, opt.output_h, opt.output_w)
# output_flip = torch.from_numpy(output_flip).cuda(
# device=opt.device, non_blocking=True)
# output_depth_flip = torch.from_numpy(output_depth_flip).cuda(
# device=opt.device, non_blocking=True)
# output[-1]['hm'] = (output[-1]['hm'] + output_flip) / 2
# output[-1]['depth'] = (output[-1]['depth'] + output_depth_flip) / 2
# pred = get_preds(output[-1]['hm'].detach().cpu().numpy())
# preds.append(convert_eval_format(pred, conf, meta)[0])
acc = accuracy_ocv(preds, torch.argmax(ocv_gt, 1))
Loss_ocv.update(loss.item(), img.size(0))
Acc_ocv.update(acc, img.size(0))
# Loss.update(loss.item(), batch['input'].size(0))
# Loss3D.update(loss_3d.item(), batch['input'].size(0))
# Acc.update(accuracy(output[-1]['hm'].detach().cpu().numpy(),
# batch['target'].detach().cpu().numpy(), acc_idxs))
# mpeje_batch, mpjpe_cnt = mpjpe(output[-1]['hm'].detach().cpu().numpy(),
# output[-1]['depth'].detach().cpu().numpy(),
# batch['meta']['gt_3d'].detach().numpy(),
# convert_func=convert_eval_format)
# MPJPE.update(mpeje_batch, mpjpe_cnt)
batch_time.update(time.time() - end)
end = time.time()
if not opt.hide_data_time:
time_str = ' |Data {dt.avg:.3f}s({dt.val:.3f}s)' \
' |Net {bt.avg:.3f}s'.format(dt=data_time, bt=batch_time)
# Bar.suffix = '{split}: [{0}][{1}/{2}] |Total {total:} |ETA {eta:} '\
# '|Loss {loss.avg:.5f} |Loss3D {loss_3d.avg:.5f}'\
# '|Acc {Acc.avg:.4f} |MPJPE {MPJPE.avg:.2f}'\
# '{time_str}'.format(epoch, i, nIters, total=bar.elapsed_td,
# eta=bar.eta_td, loss=Loss, Acc=Acc,
# split=split, time_str=time_str,
# MPJPE=MPJPE, loss_3d=Loss3D)
Bar.suffix = '{split}: [{0}][{1}/{2}] |Total {total:} |ETA {eta:} '\
'|Loss_ocv {loss.avg:.5f}'\
'|Acc_ocv {Acc.avg:.4f}'\
'|loss_batch {loss_batch:.4f}'\
'|acc_batch {acc_batch:.4f}'\
'{time_str}'.format(epoch, i, nIters, total=bar.elapsed_td,
eta=bar.eta_td, loss=Loss_ocv, Acc=Acc_ocv,
loss_batch=loss.item(), acc_batch=acc,
split=split, time_str=time_str)
if opt.print_iter > 0:
if i % opt.print_iter == 0:
print('{}| {}'.format(opt.exp_id, Bar.suffix))
else:
bar.next()
if opt.debug >= 2:
gt = get_preds(batch['target'].cpu().numpy()) * 4
pred = get_preds(output[-1]['hm'].detach().cpu().numpy()) * 4
debugger = Debugger(ipynb=opt.print_iter > 0, edges=edges)
img = (batch['input'][0].cpu().numpy().transpose(1, 2, 0) * std +
mean) * 256
img = img.astype(np.uint8).copy()
debugger.add_img(img)
debugger.add_mask(
cv2.resize(batch['target'][0].cpu().numpy().max(axis=0),
(opt.input_w, opt.input_h)), img, 'target')
debugger.add_mask(
cv2.resize(
output[-1]['hm'][0].detach().cpu().numpy().max(axis=0),
(opt.input_w, opt.input_h)), img, 'pred')
debugger.add_point_2d(gt[0], (0, 0, 255))
debugger.add_point_2d(pred[0], (255, 0, 0))
debugger.add_point_3d(batch['meta']['gt_3d'].detach().numpy()[0],
'r',
edges=edges_3d)
pred_3d = get_preds_3d(output[-1]['hm'].detach().cpu().numpy(),
output[-1]['depth'].detach().cpu().numpy())
debugger.add_point_3d(convert_eval_format(pred_3d[0]),
'b',
edges=edges_3d)
debugger.show_all_imgs(pause=False)
debugger.show_3d()
# pdb.set_trace()
bar.finish()
# return {'loss': Loss.avg,
# 'acc': Acc.avg,
# 'mpjpe': MPJPE.avg,
# 'time': bar.elapsed_td.total_seconds() / 60.}, preds
return {
'loss': Loss_ocv.avg,
'acc': Acc_ocv.avg,
'time': bar.elapsed_td.total_seconds() / 60.
}, preds
def run(self, image_or_path_or_tensor, meta=None):
load_time, pre_time, net_time, dec_time, post_time = 0, 0, 0, 0, 0
merge_time, tot_time = 0, 0
image_name = None
debugger = Debugger(dataset=self.opt.dataset,
ipynb=(self.opt.debug == 3),
theme=self.opt.debugger_theme)
start_time = time.time()
pre_processed = False
if isinstance(image_or_path_or_tensor, np.ndarray):
image = image_or_path_or_tensor
elif type(image_or_path_or_tensor) == type(''):
image = cv2.imread(image_or_path_or_tensor)
import os
image_name = os.path.basename(image_or_path_or_tensor)
else:
image = image_or_path_or_tensor['image'][0].numpy()
pre_processed_images = image_or_path_or_tensor
pre_processed = True
loaded_time = time.time()
load_time += (loaded_time - start_time)
detections = []
for scale in self.scales:
scale_start_time = time.time()
if not pre_processed:
images, meta = self.pre_process(image, scale, meta)
else:
# import pdb; pdb.set_trace()
images = pre_processed_images['images'][scale][0]
meta = pre_processed_images['meta'][scale]
meta = {k: v.numpy()[0] for k, v in meta.items()}
images = images.to(self.opt.device)
torch.cuda.synchronize()
pre_process_time = time.time()
pre_time += pre_process_time - scale_start_time
output, dets, forward_time = self.process(images, return_time=True)
print(image_name)
print(dets)
#if self.reconstruct_img:
# self.save_tensor_to_img(output['reconstruct_img'])
torch.cuda.synchronize()
net_time += forward_time - pre_process_time
decode_time = time.time()
dec_time += decode_time - forward_time
if self.opt.debug >= 2:
self.debug(debugger, images, dets, output, scale)
dets = self.post_process(dets, meta, scale)
#print(dets)
torch.cuda.synchronize()
post_process_time = time.time()
post_time += post_process_time - decode_time
detections.append(dets)
results = self.merge_outputs(detections)
torch.cuda.synchronize()
end_time = time.time()
merge_time += end_time - post_process_time
tot_time += end_time - start_time
if self.opt.debug >= 1:
if image_name is not None:
self.show_results(debugger, image, results)
#self.save_person_only(debugger, image, results, image_name)
if self.opt.debug == 0:
if image_name is not None:
self.save_results_only(debugger, image, results, image_name)
return {
'results': results,
'tot': tot_time,
'load': load_time,
'pre': pre_time,
'net': net_time,
'dec': dec_time,
'post': post_time,
'merge': merge_time,
'image_name': image_name
}
def demo(opt):
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
opt.debug = max(opt.debug, 1)
debugger = Debugger(dataset=opt.dataset,
ipynb=(opt.debug == 3),
theme=opt.debugger_theme)
model = CenterHarDNet(85, opt).cuda()
model.eval()
image_size = (opt.input_w, opt.input_h)
det_size = [
image_size[0] // opt.down_ratio, image_size[1] // opt.down_ratio
]
onnx_model_path = "ctdet_%s_%dx%d.onnx" % (opt.arch, image_size[0],
image_size[1])
trt_engine_path = "ctdet_%s_%dx%d.trt" % (opt.arch, image_size[0],
image_size[1])
x = torch.randn((1, 3, image_size[1], image_size[0])).float().cuda()
if opt.load_trt:
trt_engine_path = opt.load_trt
engine = get_engine(1, "", trt_engine_path, fp16_mode=True)
else:
if not opt.load_model:
print('Please load model with --load_model')
exit()
print('\nStep 1: Converting ONNX... (PyTorch>1.3 is required)')
torch.onnx.export(
model, # model being run
x, # model input (or a tuple for multiple inputs)
onnx_model_path, # where to save the model (can be a file or file-like object)
export_params=
True, # store the trained parameter weights inside the model file
opset_version=11, # the ONNX version to export the model to
do_constant_folding=
True, # whether to execute constant folding for optimization
input_names=['input'], # the model's input names
output_names=['output'], # the model's output names
dynamic_axes=None)
print('\nStep 2: Converting TensorRT... ')
engine = get_engine(1,
onnx_model_path,
trt_engine_path,
fp16_mode=True,
overwrite=True)
outs = model(x)
out_shape = outs.shape
context = engine.create_execution_context()
inputs, outputs, bindings, stream = allocate_buffers(engine)
if not opt.demo:
exit()
if opt.demo == 'webcam' or \
opt.demo[opt.demo.rfind('.') + 1:].lower() in video_ext:
cam = cv2.VideoCapture(0 if opt.demo == 'webcam' else opt.demo)
cam.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
cam.set(cv2.CAP_PROP_FPS, 30)
while True:
t = perf_counter()
_, img = cam.read()
img_x = cv2.resize(img, image_size)
x = img_x.transpose(2, 0,
1).reshape(1, 3, image_size[1],
image_size[0]).astype(np.float32)
inputs[0].host = x.reshape(-1)
t2 = perf_counter()
outs = do_inference(context,
bindings=bindings,
inputs=inputs,
outputs=outputs,
stream=stream)
dets = outs[0].reshape(out_shape)
t3 = perf_counter()
show_det(dets, img, det_size, debugger, opt, False)
print(' Latency = %.2f ms (net=%.2f ms)' %
((perf_counter() - t) * 1000, (t3 - t2) * 1000))
if cv2.waitKey(1) == 27:
return # esc to quit
else:
if os.path.isdir(opt.demo):
image_names = []
ls = os.listdir(opt.demo)
for file_name in sorted(ls):
ext = file_name[file_name.rfind('.') + 1:].lower()
if ext in image_ext:
image_names.append(os.path.join(opt.demo, file_name))
else:
image_names = [opt.demo]
for (image_name) in image_names:
img = cv2.imread(image_name)
img_x = cv2.resize(img, image_size)
x = img_x.transpose(2, 0,
1).reshape(1, 3, image_size[1],
image_size[0]).astype(np.float32)
inputs[0].host = x.reshape(-1)
outs = do_inference(context,
bindings=bindings,
inputs=inputs,
outputs=outputs,
stream=stream)
dets = outs[0].reshape(out_shape)
show_det(dets, img, det_size, debugger, opt, True,
os.path.basename(image_name))
def debug(self, batch, output, iter_id, dataset):
opt = self.opt
if 'pre_hm' in batch:
output.update({'pre_hm': batch['pre_hm']})
dets = generic_decode(output, K=opt.K, opt=opt)
for k in dets:
dets[k] = dets[k].detach().cpu().numpy()
dets_gt = batch['meta']['gt_det']
for i in range(1):
debugger = Debugger(opt=opt, dataset=dataset)
img = batch['image'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = np.clip(((img * dataset.std + dataset.mean) * 255.), 0,
255).astype(np.uint8)
pred = debugger.gen_colormap(
output['hm'][i].detach().cpu().numpy())
gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'pred_hm')
debugger.add_blend_img(img, gt, 'gt_hm')
if 'pre_img' in batch:
pre_img = batch['pre_img'][i].detach().cpu().numpy().transpose(
1, 2, 0)
pre_img = np.clip(
((pre_img * dataset.std + dataset.mean) * 255), 0,
255).astype(np.uint8)
debugger.add_img(pre_img, 'pre_img_pred')
debugger.add_img(pre_img, 'pre_img_gt')
if 'pre_hm' in batch:
pre_hm = debugger.gen_colormap(
batch['pre_hm'][i].detach().cpu().numpy())
debugger.add_blend_img(pre_img, pre_hm, 'pre_hm')
debugger.add_img(img, img_id='out_pred')
if 'ltrb_amodal' in opt.heads:
debugger.add_img(img, img_id='out_pred_amodal')
debugger.add_img(img, img_id='out_gt_amodal')
# Predictions
for k in range(len(dets['scores'][i])):
if dets['scores'][i, k] > opt.vis_thresh:
debugger.add_coco_bbox(dets['bboxes'][i, k] *
opt.down_ratio,
dets['clses'][i, k],
dets['scores'][i, k],
img_id='out_pred')
if 'ltrb_amodal' in opt.heads:
debugger.add_coco_bbox(dets['bboxes_amodal'][i, k] *
opt.down_ratio,
dets['clses'][i, k],
dets['scores'][i, k],
img_id='out_pred_amodal')
if 'hps' in opt.heads and int(dets['clses'][i, k]) == 0:
debugger.add_coco_hp(dets['hps'][i, k] *
opt.down_ratio,
img_id='out_pred')
if 'tracking' in opt.heads:
debugger.add_arrow(dets['cts'][i][k] * opt.down_ratio,
dets['tracking'][i][k] *
opt.down_ratio,
img_id='out_pred')
debugger.add_arrow(dets['cts'][i][k] * opt.down_ratio,
dets['tracking'][i][k] *
opt.down_ratio,
img_id='pre_img_pred')
# Ground truth
debugger.add_img(img, img_id='out_gt')
for k in range(len(dets_gt['scores'][i])):
if dets_gt['scores'][i][k] > opt.vis_thresh:
debugger.add_coco_bbox(dets_gt['bboxes'][i][k] *
opt.down_ratio,
dets_gt['clses'][i][k],
dets_gt['scores'][i][k],
img_id='out_gt')
if 'ltrb_amodal' in opt.heads:
debugger.add_coco_bbox(dets_gt['bboxes_amodal'][i, k] *
opt.down_ratio,
dets_gt['clses'][i, k],
dets_gt['scores'][i, k],
img_id='out_gt_amodal')
if 'hps' in opt.heads and \
(int(dets['clses'][i, k]) == 0):
debugger.add_coco_hp(dets_gt['hps'][i][k] *
opt.down_ratio,
img_id='out_gt')
if 'tracking' in opt.heads:
debugger.add_arrow(
dets_gt['cts'][i][k] * opt.down_ratio,
dets_gt['tracking'][i][k] * opt.down_ratio,
img_id='out_gt')
debugger.add_arrow(
dets_gt['cts'][i][k] * opt.down_ratio,
dets_gt['tracking'][i][k] * opt.down_ratio,
img_id='pre_img_gt')
if 'hm_hp' in opt.heads:
pred = debugger.gen_colormap_hp(
output['hm_hp'][i].detach().cpu().numpy())
gt = debugger.gen_colormap_hp(
batch['hm_hp'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'pred_hmhp')
debugger.add_blend_img(img, gt, 'gt_hmhp')
if 'rot' in opt.heads and 'dim' in opt.heads and 'dep' in opt.heads:
dets_gt = {k: dets_gt[k].cpu().numpy() for k in dets_gt}
calib = batch['meta']['calib'].detach().numpy() \
if 'calib' in batch['meta'] else None
det_pred = generic_post_process(
opt, dets, batch['meta']['c'].cpu().numpy(),
batch['meta']['s'].cpu().numpy(), output['hm'].shape[2],
output['hm'].shape[3], self.opt.num_classes, calib)
det_gt = generic_post_process(opt, dets_gt,
batch['meta']['c'].cpu().numpy(),
batch['meta']['s'].cpu().numpy(),
output['hm'].shape[2],
output['hm'].shape[3],
self.opt.num_classes, calib)
debugger.add_3d_detection(batch['meta']['img_path'][i],
batch['meta']['flipped'][i],
det_pred[i],
calib[i],
vis_thresh=opt.vis_thresh,
img_id='add_pred')
debugger.add_3d_detection(batch['meta']['img_path'][i],
batch['meta']['flipped'][i],
det_gt[i],
calib[i],
vis_thresh=opt.vis_thresh,
img_id='add_gt')
debugger.add_bird_views(det_pred[i],
det_gt[i],
vis_thresh=opt.vis_thresh,
img_id='bird_pred_gt')
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir,
prefix='{}'.format(iter_id))
else:
debugger.show_all_imgs(pause=True)
def debug(self, batch, output, iter_id):
opt = self.opt
wh = output['wh'] if opt.reg_bbox else None
reg = output['reg'] if opt.reg_offset else None
dets = ctadd_decode(output['hm'],
output['rot'],
output['dep'],
output['dim'],
wh=wh,
reg=reg,
K=opt.K)
# x, y, score, r1-r8, depth, dim1-dim3, cls
dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
calib = batch['meta']['calib'].detach().numpy()
# x, y, score, rot, depth, dim1, dim2, dim3
# if opt.dataset == 'gta':
# dets[:, 12:15] /= 3
dets_pred = ctadd_post_process(dets.copy(),
batch['meta']['c'].detach().numpy(),
batch['meta']['s'].detach().numpy(),
calib, opt)
dets_gt = ctadd_post_process(
batch['meta']['gt_det'].detach().numpy().copy(),
batch['meta']['c'].detach().numpy(),
batch['meta']['s'].detach().numpy(), calib, opt)
#for i in range(input.size(0)):
for i in range(1):
debugger = Debugger(dataset=opt.dataset,
ipynb=(opt.debug == 3),
theme=opt.debugger_theme)
img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = ((img * self.opt.std + self.opt.mean) * 255.).astype(
np.uint8)
pred = debugger.gen_colormap(
output['hm'][i].detach().cpu().numpy())
gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'hm_pred')
debugger.add_blend_img(img, gt, 'hm_gt')
# decode
debugger.add_ct_detection(img,
dets[i],
show_box=opt.reg_bbox,
center_thresh=opt.center_thresh,
img_id='det_pred')
debugger.add_ct_detection(
img,
batch['meta']['gt_det'][i].cpu().numpy().copy(),
show_box=opt.reg_bbox,
img_id='det_gt')
debugger.add_3d_detection(batch['meta']['image_path'][i],
dets_pred[i],
calib[i],
center_thresh=opt.center_thresh,
img_id='add_pred')
debugger.add_3d_detection(batch['meta']['image_path'][i],
dets_gt[i],
calib[i],
center_thresh=opt.center_thresh,
img_id='add_gt')
# debugger.add_bird_view(
# dets_pred[i], center_thresh=opt.center_thresh, img_id='bird_pred')
# debugger.add_bird_view(dets_gt[i], img_id='bird_gt')
debugger.add_bird_views(dets_pred[i],
dets_gt[i],
center_thresh=opt.center_thresh,
img_id='bird_pred_gt')
# debugger.add_blend_img(img, pred, 'out', white=True)
debugger.compose_vis_add(batch['meta']['image_path'][i],
dets_pred[i],
calib[i],
opt.center_thresh,
pred,
'bird_pred_gt',
img_id='out')
# debugger.add_img(img, img_id='out')
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir,
prefix='{}'.format(iter_id))
else:
debugger.show_all_imgs(pause=True)
def train(self, cfg):
# 设置gpu环境,考虑单卡多卡情况
gpus_str = ''
if isinstance(cfg.gpus, (list, tuple)):
cfg.gpus = [int(i) for i in cfg.gpus]
for s in cfg.gpus:
gpus_str += str(s) + ','
gpus_str = gpus_str[:-1]
else:
gpus_str = str(int(cfg.gpus))
cfg.gpus = [int(cfg.gpus)]
os.environ['CUDA_VISIBLE_DEVICES'] = gpus_str
cfg.gpus = [i for i in range(len(cfg.gpus))
] if cfg.gpus[0] >= 0 else [-1]
# 设置log
model_dir = os.path.join(cfg.save_dir, cfg.id)
debug_dir = os.path.join(model_dir, 'debug')
if not os.path.exists(model_dir):
os.makedirs(model_dir)
if not os.path.exists(debug_dir):
os.makedirs(debug_dir)
logger = setup_logger(cfg.id, os.path.join(model_dir, 'log'))
if USE_TENSORBOARD:
writer = tensorboardX.SummaryWriter(
log_dir=os.path.join(model_dir, 'log'))
logger.info(cfg)
gpus = cfg.gpus
device = torch.device('cpu' if gpus[0] < 0 else 'cuda')
lr = cfg.lr
lr_step = cfg.lr_step
num_epochs = cfg.num_epochs
val_step = cfg.val_step
sample_size = cfg.sample_size
# 设置数据集
dataset = YOLO(cfg.data_dir,
cfg.hflip,
cfg.vflip,
cfg.rotation,
cfg.scale,
cfg.shear,
opt=cfg,
split='train')
names = dataset.class_name
std = dataset.std
mean = dataset.mean
# 用数据集类别数设置预测网络
cfg.setup_head(dataset)
trainloader = DataLoader(dataset,
batch_size=cfg.batch_size,
shuffle=True,
num_workers=cfg.num_workers,
pin_memory=True,
drop_last=True)
# val_dataset = YOLO(cfg.data_dir, cfg.hflip, cfg.vflip, cfg.rotation, cfg.scale, cfg.shear, opt=cfg, split='val')
# valloader = DataLoader(val_dataset, batch_size=1, shuffle=True, num_workers=1, pin_memory=True)
valid_file = cfg.val_dir if not cfg.val_dir == '' else os.path.join(
cfg.data_dir, 'valid.txt')
with open(valid_file, 'r') as f:
val_list = [l.rstrip() for l in f.readlines()]
net = create_model(cfg.arch, cfg.heads, cfg.head_conv, cfg.down_ratio,
cfg.filters)
optimizer = optim.Adam(net.parameters(), lr=lr)
start_epoch = 0
if cfg.resume:
pretrain = os.path.join(model_dir, 'model_last.pth')
if os.path.exists(pretrain):
print('resume model from %s' % pretrain)
try:
net, optimizer, start_epoch = load_model(
net, pretrain, optimizer, True, lr, lr_step)
except:
print('\t... loading model error: ckpt may not compatible')
model = ModleWithLoss(net, CtdetLoss(cfg))
if len(gpus) > 1:
model = nn.DataParallel(model, device_ids=gpus).to(device)
else:
model = model.to(device)
step = 0
best = 1e10
log_loss_stats = ['loss', 'hm_loss', 'wh_loss']
if cfg.reg_offset:
log_loss_stats += ['off_loss']
if cfg.reg_obj:
log_loss_stats += ['obj_loss']
for epoch in range(start_epoch + 1, num_epochs + 1):
avg_loss_stats = {l: AverageMeter() for l in log_loss_stats}
model.train()
with tqdm(trainloader) as loader:
for _, batch in enumerate(loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=device,
non_blocking=True)
output, loss, loss_stats = model(batch)
loss = loss.mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 设置tqdm显示信息
lr = optimizer.param_groups[0]['lr']
poststr = ''
for l in avg_loss_stats:
avg_loss_stats[l].update(loss_stats[l].mean().item(),
batch['input'].size(0))
poststr += '{}: {:.4f}; '.format(
l, avg_loss_stats[l].avg)
loader.set_description('Epoch %d' % (epoch))
poststr += 'lr: {:.4f}'.format(lr)
loader.set_postfix_str(poststr)
step += 1
# self.lossSignal.emit(loss.item(), step)
del output, loss, loss_stats
# valid
if step % val_step == 0:
if len(cfg.gpus) > 1:
val_model = model.module
else:
val_model = model
val_model.eval()
torch.cuda.empty_cache()
# 随机采样
idx = np.arange(len(val_list))
idx = np.random.permutation(idx)[:sample_size]
for j, id in enumerate(idx):
image = cv2.imread(val_list[id])
image = self.preprocess(image, cfg.input_h,
cfg.input_w, mean, std)
image = image.to(device)
with torch.no_grad():
output = val_model.model(image)[-1]
# 画图并保存
debugger = Debugger(dataset=names,
down_ratio=cfg.down_ratio)
reg = output['reg'] if cfg.reg_offset else None
obj = output['obj'] if cfg.reg_obj else None
dets = ctdet_decode(output['hm'].sigmoid_(),
output['wh'],
reg=reg,
obj=obj,
cat_spec_wh=cfg.cat_spec_wh,
K=cfg.K)
dets = dets.detach().cpu().numpy().reshape(
-1, dets.shape[2])
dets[:, :4] *= cfg.down_ratio
image = image[0].detach().cpu().numpy().transpose(
1, 2, 0)
image = np.clip(((image * std + mean) * 255.), 0,
255).astype(np.uint8)
pred = debugger.gen_colormap(
output['hm'][0].detach().cpu().numpy())
debugger.add_blend_img(image, pred, 'pred_hm')
debugger.add_img(image, img_id='out_pred')
for k in range(len(dets)):
if dets[k, 4] > cfg.vis_thresh:
debugger.add_coco_bbox(dets[k, :4],
dets[k, -1],
dets[k, 4],
img_id='out_pred')
debugger.save_all_imgs(debug_dir,
prefix='{}.{}_'.format(
step, j))
del output, image, dets
# 保存模型参数
save_model(os.path.join(model_dir, 'model_best.pth'),
epoch, net)
model.train()
logstr = 'epoch {}'.format(epoch)
for k, v in avg_loss_stats.items():
logstr += ' {}: {:.4f};'.format(k, v.avg)
if USE_TENSORBOARD:
writer.add_scalar('train_{}'.format(k), v.avg, epoch)
logger.info(logstr)
# if epoch % val_step == 0:
# if len(cfg.gpus) > 1:
# val_model = model.module
# else:
# val_model = model
# val_model.eval()
# torch.cuda.empty_cache()
#
# val_loss_stats = {l: AverageMeter() for l in log_loss_stats}
#
# with tqdm(valloader) as loader:
# for j, batch in enumerate(loader):
# for k in batch:
# if k != 'meta':
# batch[k] = batch[k].to(device=device, non_blocking=True)
# with torch.no_grad():
# output, loss, loss_stats = val_model(batch)
#
# poststr = ''
# for l in val_loss_stats:
# val_loss_stats[l].update(
# loss_stats[l].mean().item(), batch['input'].size(0))
# poststr += '{}: {:.4f}; '.format(l, val_loss_stats[l].avg)
# loader.set_description('Epoch %d valid' % (epoch))
# poststr += 'lr: {:.4f}'.format(lr)
# loader.set_postfix_str(poststr)
#
# if j < sample_size:
# # 将预测结果画出来保存成jpg图片
# debugger = Debugger(dataset=names, down_ratio=cfg.down_ratio)
# reg = output['reg'] if cfg.reg_offset else None
# obj = output['obj'] if cfg.reg_obj else None
# dets = ctdet_decode(
# output['hm'], output['wh'], reg=reg, obj=obj,
# cat_spec_wh=cfg.cat_spec_wh, K=cfg.K)
# dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
# dets[:, :, :4] *= cfg.down_ratio
# dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])
# dets_gt[:, :, :4] *= cfg.down_ratio
# for i in range(1):
# img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
# img = np.clip(((img * std + mean) * 255.), 0, 255).astype(np.uint8)
# pred = debugger.gen_colormap(output['hm'][i].detach().cpu().numpy())
# gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
# debugger.add_blend_img(img, pred, 'pred_hm')
# debugger.add_blend_img(img, gt, 'gt_hm')
# debugger.add_img(img, img_id='out_pred')
# for k in range(len(dets[i])):
# if dets[i, k, 4] > cfg.vis_thresh:
# debugger.add_coco_bbox(dets[i, k, :4], dets[i, k, -1],
# dets[i, k, 4], img_id='out_pred')
#
# debugger.add_img(img, img_id='out_gt')
# for k in range(len(dets_gt[i])):
# if dets_gt[i, k, 4] > cfg.vis_thresh:
# debugger.add_coco_bbox(dets_gt[i, k, :4], dets_gt[i, k, -1],
# dets_gt[i, k, 4], img_id='out_gt')
#
# debugger.save_all_imgs(debug_dir, prefix='{}.{}_'.format(epoch, j))
# del output, loss, loss_stats
# model.train()
# logstr = 'epoch {} valid'.format(epoch)
# for k, v in val_loss_stats.items():
# logstr += ' {}: {:.4f};'.format(k, v.avg)
# if USE_TENSORBOARD:
# writer.add_scalar('val_{}'.format(k), v.avg, epoch)
# logger.info(logstr)
# if val_loss_stats['loss'].avg < best:
# best = val_loss_stats['loss'].avg
# save_model(os.path.join(model_dir, 'model_best.pth'), epoch, net)
save_model(os.path.join(model_dir, 'model_last.pth'), epoch, net,
optimizer)
if epoch in cfg.lr_step:
save_model(
os.path.join(model_dir, 'model_{}.pth'.format(epoch)),
epoch, net, optimizer)
lr = cfg.lr * (0.1**(cfg.lr_step.index(epoch) + 1))
logger.info('Drop LR to {}'.format(lr))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def debug(self, batch, output, iter_id, dataset):
opt = self.opt
if 'pre_hm' in batch:
output.update({'pre_hm': batch['pre_hm']})
dets = fusion_decode(output, K=opt.K, opt=opt)
for k in dets:
dets[k] = dets[k].detach().cpu().numpy()
dets_gt = batch['meta']['gt_det']
for i in range(1):
debugger = Debugger(opt=opt, dataset=dataset)
img = batch['image'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = np.clip(((img * dataset.std + dataset.mean) * 255.), 0,
255).astype(np.uint8)
pred = debugger.gen_colormap(
output['hm'][i].detach().cpu().numpy())
gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
debugger.add_blend_img(img,
pred,
'pred_hm',
trans=self.opt.hm_transparency)
debugger.add_blend_img(img,
gt,
'gt_hm',
trans=self.opt.hm_transparency)
debugger.add_img(img, img_id='img')
# show point clouds
if opt.pointcloud:
pc_2d = batch['pc_2d'][i].detach().cpu().numpy()
pc_3d = None
pc_N = batch['pc_N'][i].detach().cpu().numpy()
debugger.add_img(img, img_id='pc')
debugger.add_pointcloud(pc_2d, pc_N, img_id='pc')
if 'pc_hm' in opt.pc_feat_lvl:
channel = opt.pc_feat_channels['pc_hm']
pc_hm = debugger.gen_colormap(
batch['pc_hm'][i][channel].unsqueeze(
0).detach().cpu().numpy())
debugger.add_blend_img(img,
pc_hm,
'pc_hm',
trans=self.opt.hm_transparency)
if 'pc_dep' in opt.pc_feat_lvl:
channel = opt.pc_feat_channels['pc_dep']
pc_hm = batch['pc_hm'][i][channel].unsqueeze(
0).detach().cpu().numpy()
pc_dep = debugger.add_overlay_img(img, pc_hm, 'pc_dep')
if 'pre_img' in batch:
pre_img = batch['pre_img'][i].detach().cpu().numpy().transpose(
1, 2, 0)
pre_img = np.clip(
((pre_img * dataset.std + dataset.mean) * 255), 0,
255).astype(np.uint8)
debugger.add_img(pre_img, 'pre_img_pred')
debugger.add_img(pre_img, 'pre_img_gt')
if 'pre_hm' in batch:
pre_hm = debugger.gen_colormap(
batch['pre_hm'][i].detach().cpu().numpy())
debugger.add_blend_img(pre_img,
pre_hm,
'pre_hm',
trans=self.opt.hm_transparency)
debugger.add_img(img, img_id='out_pred')
if 'ltrb_amodal' in opt.heads:
debugger.add_img(img, img_id='out_pred_amodal')
debugger.add_img(img, img_id='out_gt_amodal')
# Predictions
for k in range(len(dets['scores'][i])):
if dets['scores'][i, k] > opt.vis_thresh:
debugger.add_coco_bbox(dets['bboxes'][i, k] *
opt.down_ratio,
dets['clses'][i, k],
dets['scores'][i, k],
img_id='out_pred')
if 'ltrb_amodal' in opt.heads:
debugger.add_coco_bbox(dets['bboxes_amodal'][i, k] *
opt.down_ratio,
dets['clses'][i, k],
dets['scores'][i, k],
img_id='out_pred_amodal')
if 'hps' in opt.heads and int(dets['clses'][i, k]) == 0:
debugger.add_coco_hp(dets['hps'][i, k] *
opt.down_ratio,
img_id='out_pred')
if 'tracking' in opt.heads:
debugger.add_arrow(dets['cts'][i][k] * opt.down_ratio,
dets['tracking'][i][k] *
opt.down_ratio,
img_id='out_pred')
debugger.add_arrow(dets['cts'][i][k] * opt.down_ratio,
dets['tracking'][i][k] *
opt.down_ratio,
img_id='pre_img_pred')
# Ground truth
debugger.add_img(img, img_id='out_gt')
for k in range(len(dets_gt['scores'][i])):
if dets_gt['scores'][i][k] > opt.vis_thresh:
if 'dep' in dets_gt.keys():
dist = dets_gt['dep'][i][k]
if len(dist) > 1:
dist = dist[0]
else:
dist = -1
debugger.add_coco_bbox(dets_gt['bboxes'][i][k] *
opt.down_ratio,
dets_gt['clses'][i][k],
dets_gt['scores'][i][k],
img_id='out_gt',
dist=dist)
if 'ltrb_amodal' in opt.heads:
debugger.add_coco_bbox(dets_gt['bboxes_amodal'][i, k] *
opt.down_ratio,
dets_gt['clses'][i, k],
dets_gt['scores'][i, k],
img_id='out_gt_amodal')
if 'hps' in opt.heads and \
(int(dets['clses'][i, k]) == 0):
debugger.add_coco_hp(dets_gt['hps'][i][k] *
opt.down_ratio,
img_id='out_gt')
if 'tracking' in opt.heads:
debugger.add_arrow(
dets_gt['cts'][i][k] * opt.down_ratio,
dets_gt['tracking'][i][k] * opt.down_ratio,
img_id='out_gt')
debugger.add_arrow(
dets_gt['cts'][i][k] * opt.down_ratio,
dets_gt['tracking'][i][k] * opt.down_ratio,
img_id='pre_img_gt')
if 'hm_hp' in opt.heads:
pred = debugger.gen_colormap_hp(
output['hm_hp'][i].detach().cpu().numpy())
gt = debugger.gen_colormap_hp(
batch['hm_hp'][i].detach().cpu().numpy())
debugger.add_blend_img(img,
pred,
'pred_hmhp',
trans=self.opt.hm_transparency)
debugger.add_blend_img(img,
gt,
'gt_hmhp',
trans=self.opt.hm_transparency)
if 'rot' in opt.heads and 'dim' in opt.heads and 'dep' in opt.heads:
dets_gt = {k: dets_gt[k].cpu().numpy() for k in dets_gt}
calib = batch['meta']['calib'].detach().numpy() \
if 'calib' in batch['meta'] else None
det_pred = generic_post_process(
opt, dets, batch['meta']['c'].cpu().numpy(),
batch['meta']['s'].cpu().numpy(), output['hm'].shape[2],
output['hm'].shape[3], self.opt.num_classes, calib)
det_gt = generic_post_process(opt,
dets_gt,
batch['meta']['c'].cpu().numpy(),
batch['meta']['s'].cpu().numpy(),
output['hm'].shape[2],
output['hm'].shape[3],
self.opt.num_classes,
calib,
is_gt=True)
debugger.add_3d_detection(batch['meta']['img_path'][i],
batch['meta']['flipped'][i],
det_pred[i],
calib[i],
vis_thresh=opt.vis_thresh,
img_id='add_pred')
debugger.add_3d_detection(batch['meta']['img_path'][i],
batch['meta']['flipped'][i],
det_gt[i],
calib[i],
vis_thresh=opt.vis_thresh,
img_id='add_gt')
pc_3d = None
if opt.pointcloud:
pc_3d = batch['pc_3d'].cpu().numpy()
debugger.add_bird_views(det_pred[i],
det_gt[i],
vis_thresh=opt.vis_thresh,
img_id='bird_pred_gt',
pc_3d=pc_3d,
show_velocity=opt.show_velocity)
debugger.add_bird_views([],
det_gt[i],
vis_thresh=opt.vis_thresh,
img_id='bird_gt',
pc_3d=pc_3d,
show_velocity=opt.show_velocity)
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir,
prefix='{}'.format(iter_id))
else:
debugger.show_all_imgs(pause=True)
def step(split, epoch, opt, dataLoader, model, criterion, optimizer=None):
if split == 'train':
model.train()
else:
model.eval()
Loss, Acc, Mpjpe, Loss3D = AverageMeter(), AverageMeter(), AverageMeter(
), AverageMeter()
nIters = len(dataLoader)
bar = Bar('==>', max=nIters)
for i, (input, target2D, target3D, meta) in enumerate(dataLoader):
input_var = torch.autograd.Variable(input).float().cuda()
target2D_var = torch.autograd.Variable(target2D).float().cuda()
target3D_var = torch.autograd.Variable(target3D).float().cuda()
output = model(input_var)
reg = output[opt.nStack]
if opt.DEBUG >= 2:
gt = getPreds(target2D.cpu().numpy()) * 4
pred = getPreds((output[opt.nStack - 1].data).cpu().numpy()) * 4
debugger = Debugger()
debugger.addImg(
(input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8))
debugger.addPoint2D(pred[0], (255, 0, 0))
debugger.addPoint2D(gt[0], (0, 0, 255))
debugger.showImg()
debugger.saveImg('debug/{}.png'.format(i))
loss = FusionCriterion(opt.regWeight, opt.varWeight)(reg, target3D_var)
Loss3D.update(loss.data[0], input.size(0))
for k in range(opt.nStack):
loss += criterion(output[k], target2D_var)
Loss.update(loss.data[0], input.size(0))
Acc.update(
Accuracy((output[opt.nStack - 1].data).cpu().numpy(),
(target2D_var.data).cpu().numpy()))
mpjpe, num3D = MPJPE((output[opt.nStack - 1].data).cpu().numpy(),
(reg.data).cpu().numpy(), meta)
if num3D > 0:
Mpjpe.update(mpjpe, num3D)
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
Bar.suffix = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | Loss3D {loss3d.avg:.6f} | Acc {Acc.avg:.6f} | Mpjpe {Mpjpe.avg:.6f} ({Mpjpe.val:.6f})'.format(
epoch,
i,
nIters,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=Loss,
Acc=Acc,
split=split,
Mpjpe=Mpjpe,
loss3d=Loss3D)
bar.next()
bar.finish()
return Loss.avg, Acc.avg, Mpjpe.avg, Loss3D.avg
def run(self, image_or_path_or_tensor, meta=None):
load_time, pre_time, net_time, dec_time, post_time = 0, 0, 0, 0, 0
merge_time, tot_time = 0, 0
debugger = Debugger(
dataset=self.opt.dataset,
ipynb=(self.opt.debug == 3),
theme=self.opt.debugger_theme # white
)
start_time = time.time()
pre_processed = False
img_name = None
if isinstance(image_or_path_or_tensor, np.ndarray): # cv read
image = image_or_path_or_tensor
elif isinstance(image_or_path_or_tensor, str): # img name
image = cv2.imread(image_or_path_or_tensor)
img_name = os.path.basename(image_or_path_or_tensor)
else: # pre_processed_images = {}
image = image_or_path_or_tensor['image'][0].numpy() # prefetch_test, PrefetchDataset
# print(image.shape) # (1080, 1920, 3), ori img size
pre_processed_images = image_or_path_or_tensor
pre_processed = True
loaded_time = time.time()
load_time += (loaded_time - start_time)
detections = []
for scale in self.scales: # default [1.0]
scale_start_time = time.time()
# pre_process input images
# images: may flip test
# meta: recover res to ori size
if not pre_processed:
# totensor, affine meta params
images, meta = self.pre_process(image, scale, meta)
else:
# import pdb; pdb.set_trace()
images = pre_processed_images['images'][scale][0] # [1, 3, 352, 640]
meta = pre_processed_images['meta'][scale]
meta = {k: v.numpy()[0] for k, v in meta.items()}
images = images.to(self.opt.device) # tensor to device
torch.cuda.synchronize() # sync and cal time
pre_process_time = time.time()
pre_time += pre_process_time - scale_start_time
# model forward, 1/4 results
output, dets, forward_time = self.process(images, return_time=True) # return forward time
torch.cuda.synchronize()
net_time += forward_time - pre_process_time
decode_time = time.time()
dec_time += decode_time - forward_time
if self.opt.debug >= 2:
# debug img, add pred_blend_img, x4
self.debug(debugger, images, dets, output, scale, img_name)
# affine transform
dets = self.post_process(dets, meta, scale)
torch.cuda.synchronize()
post_process_time = time.time()
post_time += post_process_time - decode_time
detections.append(dets)
# merge_outputs, use soft nms
results = self.merge_outputs(detections)
torch.cuda.synchronize()
end_time = time.time()
merge_time += end_time - post_process_time
tot_time += end_time - start_time
if self.opt.debug >= 1:
# pass one more param: img_name, to save img
self.show_results(debugger, image, results, img_name)
return {
'results': results,
'tot': tot_time,
'load': load_time,
'pre': pre_time,
'net': net_time,
'dec': dec_time,
'post': post_time,
'merge': merge_time
}
def run(self, image_or_path_or_tensor, meta=None):
load_time, pre_time, net_time, dec_time, post_time = 0, 0, 0, 0, 0
merge_time, tot_time = 0, 0
###实例化debugger对象
debugger = Debugger(dataset=self.opt.dataset,
ipynb=(self.opt.debug == 3),
theme=self.opt.debugger_theme)
start_time = time.time()
pre_processed = False
#####isinstance判断前者是否为后者的一个对象
if isinstance(image_or_path_or_tensor, np.ndarray):
image = image_or_path_or_tensor
elif type(image_or_path_or_tensor) == type(''):
image = cv2.imread(image_or_path_or_tensor)
else:
image = image_or_path_or_tensor['image'][0].numpy()
pre_processed_images = image_or_path_or_tensor
pre_processed = True
loaded_time = time.time()
load_time += (loaded_time - start_time)
detections = []
###对于各种图像的scale循环
for scale in self.scales:
scale_start_time = time.time()
if not pre_processed:
images, meta = self.pre_process(image, scale, meta)
else:
# import pdb; pdb.set_trace()
images = pre_processed_images['images'][scale][0]
meta = pre_processed_images['meta'][scale]
meta = {k: v.numpy()[0] for k, v in meta.items()}
images = images.to(self.opt.device)
torch.cuda.synchronize()
pre_process_time = time.time()
pre_time += pre_process_time - scale_start_time
output, dets, forward_time = self.process(images, return_time=True)
torch.cuda.synchronize()
net_time += forward_time - pre_process_time
decode_time = time.time()
dec_time += decode_time - forward_time
if self.opt.debug >= 2:
self.debug(debugger, images, dets, output, scale)
dets = self.post_process(dets, meta, scale)
torch.cuda.synchronize()
post_process_time = time.time()
post_time += post_process_time - decode_time
detections.append(dets)
results = self.merge_outputs(detections)
torch.cuda.synchronize()
end_time = time.time()
merge_time += end_time - post_process_time
tot_time += end_time - start_time
if self.opt.debug >= 1:
self.show_results(debugger, image, results)
return {
'results': results,
'tot': tot_time,
'load': load_time,
'pre': pre_time,
'net': net_time,
'dec': dec_time,
'post': post_time,
'merge': merge_time
}
def debug(self, batch, output, iter_id):
opt = self.opt
reg = output['reg'] if opt.reg_offset else None
dets = ctdet_decode(output['hm'],
output['wh'],
reg=reg,
cat_spec_wh=opt.cat_spec_wh,
K=opt.K)
dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
dets[:, :, :4] *= opt.down_ratio
dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])
dets_gt[:, :, :4] *= opt.down_ratio
if opt.task == 'ctdet_semseg':
seg_gt = batch['seg'][0][0].cpu().numpy()
seg_pred = output['seg'].max(1)[1].squeeze_(1).squeeze_(
0).cpu().numpy()
for i in range(1):
debugger = Debugger(opt,
dataset=opt.dataset,
ipynb=(opt.debug == 3),
theme=opt.debugger_theme)
img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
img = np.clip(((img * opt.std + opt.mean) * 255.), 0,
255).astype(np.uint8)
debugger.add_img(img, img_id='out_pred')
for k in range(len(dets[i])):
if dets[i, k, 4] > opt.vis_thresh:
debugger.add_coco_bbox(dets[i, k, :4],
dets[i, k, -1],
dets[i, k, 4],
img_id='out_pred')
debugger.add_img(img, img_id='out_gt')
for k in range(len(dets_gt[i])):
if dets_gt[i, k, 4] > opt.vis_thresh:
debugger.add_coco_bbox(dets_gt[i, k, :4],
dets_gt[i, k, -1],
dets_gt[i, k, 4],
img_id='out_gt')
if opt.save_video and opt.debug <= 1: # only save the predicted and gt images
return debugger.imgs['out_pred'], debugger.imgs[
'out_gt'] # , debugger.imgs['pred_hm'], debugger.imgs['gt_hm']
pred = debugger.gen_colormap(
output['hm'][i].detach().cpu().numpy())
gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
debugger.add_blend_img(img, pred, 'pred_hm')
debugger.add_blend_img(img, gt, 'gt_hm')
if opt.task == 'ctdet_semseg':
debugger.visualize_masks(seg_gt, img_id='out_mask_gt')
debugger.visualize_masks(seg_pred, img_id='out_mask_pred')
if opt.debug == 4:
debugger.save_all_imgs(opt.debug_dir, prefix=iter_id)
import pdb
pdb.set_trace()
if opt.save_video:
return debugger.imgs['out_pred'], debugger.imgs['out_gt']