Beispiel #1
0
    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)
Beispiel #2
0
    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:39] *= 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:39] *= 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:39], 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:39], 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)
Beispiel #3
0
    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 = ddd_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 = ddd_post_process(dets.copy(),
                                     batch['meta']['c'].detach().numpy(),
                                     batch['meta']['s'].detach().numpy(),
                                     calib, opt)
        dets_gt = ddd_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 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)

    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()
    data_time, batch_time = AverageMeter(), AverageMeter()
    preds = []
    time_str = ''

    nIters = len(data_loader)
    bar = Bar('{}'.format(opt.exp_id), max=nIters)

    end = time.time()
    for i, batch in enumerate(data_loader):
        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
        output = model(batch['input'])

        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

        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, amb_idx = get_preds(output[-1]['hm'].detach().cpu().numpy())
            # preds.append(convert_eval_format(pred, conf, meta)[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)
        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, amb_idx = get_preds(batch['target'].cpu().numpy())
            gt *= 4
            pred, amb_idx = get_preds(output[-1]['hm'].detach().cpu().numpy())
            pred *= 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, ignore_idx = get_preds_3d(
                output[-1]['hm'].detach().cpu().numpy(),
                output[-1]['depth'].detach().cpu().numpy(), amb_idx)
            debugger.add_point_3d(convert_eval_format(pred_3d[0]),
                                  'b',
                                  edges=edges_3d)
            debugger.show_all_imgs(pause=False)
            debugger.show_3d()

    bar.finish()
    return {
        'loss': Loss.avg,
        'acc': Acc.avg,
        'mpjpe': MPJPE.avg,
        'time': bar.elapsed_td.total_seconds() / 60.
    }, preds
Beispiel #5
0
import numpy as np

from opts import opts
from datasets.dataset.yolo import YOLO
from utils.debugger import Debugger

if __name__ == '__main__':
    opt = opts().parse()
    dataset = YOLO(opt.data_dir, opt.flip, opt.vflip, opt.rotate, opt.scale,
                   opt.shear, opt, 'train')
    opt = opts().update_dataset_info_and_set_heads(opt, dataset)
    for i in range(len(dataset)):
        debugger = Debugger(dataset=opt.names)
        data = dataset[i]
        img = data['input'].transpose(1, 2, 0)
        hm = data['hm']
        dets_gt = data['meta']['gt_det']
        dets_gt[:, :4] *= opt.down_ratio
        img = np.clip(((img * dataset.std + dataset.mean) * 255.), 0,
                      255).astype(np.uint8)
        pred = debugger.gen_colormap(hm)
        debugger.add_blend_img(img, pred, 'pred_hm')
        debugger.add_img(img, img_id='out_pred')
        for k in range(len(dets_gt)):
            debugger.add_coco_bbox(dets_gt[k, :4],
                                   dets_gt[k, -1],
                                   dets_gt[k, 4],
                                   img_id='out_pred')
        debugger.show_all_imgs(pause=True)
Beispiel #6
0
def step(split, epoch, opt, data_loader, model, optimizer=None):
    if split == 'train':
        model.train()
    else:
        model.eval()

    crit = torch.nn.MSELoss()

    acc_idxs = data_loader.dataset.acc_idxs
    edges = data_loader.dataset.edges
    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, Acc = AverageMeter(), AverageMeter()
    data_time, batch_time = AverageMeter(), AverageMeter()
    preds = []

    nIters = len(data_loader)
    bar = Bar('{}'.format(opt.exp_id), max=nIters)

    end = time.time()
    for i, batch in enumerate(data_loader):
        data_time.update(time.time() - end)
        input, target, meta = batch['input'], batch['target'], batch['meta']
        input_var = input.cuda(device=opt.device, non_blocking=True)
        target_var = target.cuda(device=opt.device, non_blocking=True)

        output = model(input_var)

        loss = crit(output[-1]['hm'], target_var)
        for k in range(opt.num_stacks - 1):
            loss += crit(output[k], target_var)

        if split == 'train':
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        else:
            input_ = 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_ = (output[-1].detach().cpu().numpy() + output_flip) / 2
            output_flip = torch.from_numpy(output_flip).cuda(device=opt.device,
                                                             non_blocking=True)
            output[-1]['hm'] = (output[-1]['hm'] + output_flip) / 2
            pred, conf = get_preds(output[-1]['hm'].detach().cpu().numpy(),
                                   True)
            preds.append(convert_eval_format(pred, conf, meta)[0])

        Loss.update(loss.detach()[0], input.size(0))
        Acc.update(
            accuracy(output[-1]['hm'].detach().cpu().numpy(),
                     target_var.detach().cpu().numpy(), acc_idxs))

        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)
        else:
            time_str = ''
        Bar.suffix = '{split}: [{0}][{1}/{2}] |Total {total:} |ETA {eta:}' \
                     '|Loss {loss.avg:.5f} |Acc {Acc.avg:.4f}'\
                     '{time_str}'.format(epoch, i, nIters, total=bar.elapsed_td,
                                         eta=bar.eta_td, loss=Loss, Acc=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(target.cpu().numpy()) * 4
            pred = get_preds(output[-1]['hm'].detach().cpu().numpy()) * 4
            debugger = Debugger(ipynb=opt.print_iter > 0, edges=edges)
            img = (input[0].numpy().transpose(1, 2, 0) * std + mean) * 256
            img = img.astype(np.uint8).copy()
            debugger.add_img(img)
            debugger.add_mask(
                cv2.resize(target[0].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(pred[0], (255, 0, 0))
            debugger.add_point_2d(gt[0], (0, 0, 255))
            debugger.show_all_imgs(pause=True)

    bar.finish()
    return {
        'loss': Loss.avg,
        'acc': Acc.avg,
        'time': bar.elapsed_td.total_seconds() / 60.
    }, preds
Beispiel #7
0
    def debug(self, batch, output, iter_id):
        cfg = self.cfg
        reg = output[3] if cfg.LOSS.REG_OFFSET else None
        hm_hp = output[4] if cfg.LOSS.HM_HP else None
        hp_offset = output[5] if cfg.LOSS.REG_HP_OFFSET else None
        dets = multi_pose_decode(output[0],
                                 output[1],
                                 output[2],
                                 reg=reg,
                                 hm_hp=hm_hp,
                                 hp_offset=hp_offset,
                                 K=cfg.TEST.TOPK)
        dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])

        dets[:, :, :4] *= cfg.MODEL.INPUT_RES / cfg.MODEL.OUTPUT_RES
        dets[:, :, 5:39] *= cfg.MODEL.INPUT_RES / cfg.MODEL.OUTPUT_RES
        dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])
        dets_gt[:, :, :4] *= cfg.MODEL.INPUT_RES / cfg.MODEL.OUTPUT_RES
        dets_gt[:, :, 5:39] *= cfg.MODEL.INPUT_RES / cfg.MODEL.OUTPUT_RES
        for i in range(1):
            debugger = Debugger(dataset=cfg.SAMPLE_METHOD,
                                ipynb=(cfg.DEBUG == 3),
                                theme=cfg.DEBUG_THEME)
            img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
            img = np.clip(((img * np.array(cfg.DATASET.STD).reshape(
                1, 1, 3).astype(np.float32) + cfg.DATASET.MEAN) * 255.), 0,
                          255).astype(np.uint8)
            pred = debugger.gen_colormap(output[0][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.MODEL.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:39], 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.MODEL.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:39], img_id='out_gt')

            if cfg.LOSS.HM_HP:
                pred = debugger.gen_colormap_hp(
                    output[4][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 cfg.DEBUG == 4:
                debugger.save_all_imgs(cfg.LOG_DIR,
                                       prefix='{}'.format(iter_id))
            else:
                debugger.show_all_imgs(pause=True)
Beispiel #8
0
    def debug(self, batch, output, iter_id):
        opt = self.opt
        reg = output['reg'] if opt.reg_offset else None
        # print(output)
        dets = circledet_decode(output['hm'],
                                output['cl'],
                                reg=reg,
                                cat_spec_wh=opt.cat_spec_wh,
                                K=opt.K)

        # print(dets)
        if opt.filter_boarder:
            output_h = self.opt.default_resolution[
                0] // self.opt.down_ratio  #hard coded
            output_w = self.opt.default_resolution[
                1] // self.opt.down_ratio  #hard coded
            for i in range(dets.shape[1]):
                cp = [0, 0]
                cp[0] = dets[0, i, 0]
                cp[1] = dets[0, i, 1]
                cr = dets[0, i, 2]
                if cp[0] - cr < 0 or cp[0] + cr > output_w:
                    dets[0, i, 3] = 0
                    continue
                if cp[1] - cr < 0 or cp[1] + cr > output_h:
                    dets[0, i, 3] = 0
                    continue

        dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
        dets[:, :, :3] *= opt.down_ratio
        dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])
        dets_gt[:, :, :3] *= 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])):
                # print('risk = %f' % dets[i, k, 3])
                if dets[i, k, 3] > opt.center_thresh:
                    debugger.add_coco_circle(dets[i, k, :3],
                                             dets[i, k, -1],
                                             dets[i, k, 3],
                                             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, 3] > opt.center_thresh:
                    debugger.add_coco_circle(dets_gt[i, k, :3],
                                             dets_gt[i, k, -1],
                                             dets_gt[i, k, 3],
                                             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)
Beispiel #9
0
def kp_detection(db, k_ind, data_aug, debug):
    data_rng = system_configs.data_rng
    batch_size = system_configs.batch_size

    categories = db.configs["categories"]
    input_size = db.configs["input_size"]
    output_size = db.configs["output_sizes"][0]

    border = db.configs["border"]
    lighting = db.configs["lighting"]
    rand_crop = db.configs["rand_crop"]
    rand_color = db.configs["rand_color"]
    rand_scales = db.configs["rand_scales"]
    gaussian_bump = db.configs["gaussian_bump"]
    gaussian_iou = db.configs["gaussian_iou"]
    gaussian_rad = db.configs["gaussian_radius"]

    max_tag_len = 128

    # allocating memory
    images = np.zeros((batch_size, 3, input_size[0], input_size[1]),
                      dtype=np.float32)
    t_heatmaps = np.zeros(
        (batch_size, categories, output_size[0], output_size[1]),
        dtype=np.float32)
    l_heatmaps = np.zeros(
        (batch_size, categories, output_size[0], output_size[1]),
        dtype=np.float32)
    b_heatmaps = np.zeros(
        (batch_size, categories, output_size[0], output_size[1]),
        dtype=np.float32)
    r_heatmaps = np.zeros(
        (batch_size, categories, output_size[0], output_size[1]),
        dtype=np.float32)
    ct_heatmaps = np.zeros(
        (batch_size, categories, output_size[0], output_size[1]),
        dtype=np.float32)
    t_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32)
    l_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32)
    b_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32)
    r_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32)
    t_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64)
    l_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64)
    b_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64)
    r_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64)
    ct_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64)
    tag_masks = np.zeros((batch_size, max_tag_len), dtype=np.uint8)
    tag_lens = np.zeros((batch_size, ), dtype=np.int32)

    db_size = db.db_inds.size
    for b_ind in range(batch_size):
        if not debug and k_ind == 0:
            db.shuffle_inds()

        db_ind = db.db_inds[k_ind]
        k_ind = (k_ind + 1) % db_size

        # reading image
        image_file = db.image_file(db_ind)
        image = cv2.imread(image_file)

        # reading detections
        detections, extreme_pts = db.detections(db_ind)

        # cropping an image randomly
        if rand_crop:
            image, detections, extreme_pts = random_crop_pts(image,
                                                             detections,
                                                             extreme_pts,
                                                             rand_scales,
                                                             input_size,
                                                             border=border)
        else:
            assert 0
            # image, detections = _full_image_crop(image, detections)

        image, detections, extreme_pts = _resize_image_pts(
            image, detections, extreme_pts, input_size)
        detections, extreme_pts = _clip_detections_pts(image, detections,
                                                       extreme_pts)

        width_ratio = output_size[1] / input_size[1]
        height_ratio = output_size[0] / input_size[0]

        # flipping an image randomly
        if np.random.uniform() > 0.5:
            image[:] = image[:, ::-1, :]
            width = image.shape[1]
            detections[:, [0, 2]] = width - detections[:, [2, 0]] - 1
            extreme_pts[:, :, 0] = width - extreme_pts[:, :, 0] - 1
            extreme_pts[:, 1, :], extreme_pts[:, 3, :] = \
                extreme_pts[:, 3, :].copy(), extreme_pts[:, 1, :].copy()

        image = image.astype(np.float32) / 255.
        if not debug:
            if rand_color:
                color_jittering_(data_rng, image)
                if lighting:
                    lighting_(data_rng, image, 0.1, db.eig_val, db.eig_vec)
        normalize_(image, db.mean, db.std)
        images[b_ind] = image.transpose((2, 0, 1))

        for ind, detection in enumerate(detections):
            category = int(detection[-1]) - 1
            extreme_pt = extreme_pts[ind]

            xt, yt = extreme_pt[0, 0], extreme_pt[0, 1]
            xl, yl = extreme_pt[1, 0], extreme_pt[1, 1]
            xb, yb = extreme_pt[2, 0], extreme_pt[2, 1]
            xr, yr = extreme_pt[3, 0], extreme_pt[3, 1]
            xct = (xl + xr) / 2
            yct = (yt + yb) / 2

            fxt = (xt * width_ratio)
            fyt = (yt * height_ratio)
            fxl = (xl * width_ratio)
            fyl = (yl * height_ratio)
            fxb = (xb * width_ratio)
            fyb = (yb * height_ratio)
            fxr = (xr * width_ratio)
            fyr = (yr * height_ratio)
            fxct = (xct * width_ratio)
            fyct = (yct * height_ratio)

            xt = int(fxt)
            yt = int(fyt)
            xl = int(fxl)
            yl = int(fyl)
            xb = int(fxb)
            yb = int(fyb)
            xr = int(fxr)
            yr = int(fyr)
            xct = int(fxct)
            yct = int(fyct)

            if gaussian_bump:
                width = detection[2] - detection[0]
                height = detection[3] - detection[1]

                width = math.ceil(width * width_ratio)
                height = math.ceil(height * height_ratio)

                if gaussian_rad == -1:
                    radius = gaussian_radius((height, width), gaussian_iou)
                    radius = max(0, int(radius))
                else:
                    radius = gaussian_rad
                draw_gaussian(t_heatmaps[b_ind, category], [xt, yt], radius)
                draw_gaussian(l_heatmaps[b_ind, category], [xl, yl], radius)
                draw_gaussian(b_heatmaps[b_ind, category], [xb, yb], radius)
                draw_gaussian(r_heatmaps[b_ind, category], [xr, yr], radius)
                draw_gaussian(ct_heatmaps[b_ind, category], [xct, yct], radius)
            else:
                t_heatmaps[b_ind, category, yt, xt] = 1
                l_heatmaps[b_ind, category, yl, xl] = 1
                b_heatmaps[b_ind, category, yb, xb] = 1
                r_heatmaps[b_ind, category, yr, xr] = 1

            tag_ind = tag_lens[b_ind]
            t_regrs[b_ind, tag_ind, :] = [fxt - xt, fyt - yt]
            l_regrs[b_ind, tag_ind, :] = [fxl - xl, fyl - yl]
            b_regrs[b_ind, tag_ind, :] = [fxb - xb, fyb - yb]
            r_regrs[b_ind, tag_ind, :] = [fxr - xr, fyr - yr]
            t_tags[b_ind, tag_ind] = yt * output_size[1] + xt
            l_tags[b_ind, tag_ind] = yl * output_size[1] + xl
            b_tags[b_ind, tag_ind] = yb * output_size[1] + xb
            r_tags[b_ind, tag_ind] = yr * output_size[1] + xr
            ct_tags[b_ind, tag_ind] = yct * output_size[1] + xct
            tag_lens[b_ind] += 1

    for b_ind in range(batch_size):
        tag_len = tag_lens[b_ind]
        tag_masks[b_ind, :tag_len] = 1

    if debug:
        debugger = Debugger(num_classes=1)
        t_hm = debugger.gen_colormap(t_heatmaps[0])
        l_hm = debugger.gen_colormap(l_heatmaps[0])
        b_hm = debugger.gen_colormap(b_heatmaps[0])
        r_hm = debugger.gen_colormap(r_heatmaps[0])
        ct_hm = debugger.gen_colormap(ct_heatmaps[0])
        img = images[0] * db.std.reshape(3, 1, 1) + db.mean.reshape(3, 1, 1)
        img = (img * 255).astype(np.uint8).transpose(1, 2, 0)
        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, np.maximum(np.maximum(t_hm, l_hm), np.maximum(b_hm, r_hm)),
            'extreme')
        debugger.add_blend_img(img, ct_hm, 'center')
        debugger.show_all_imgs(pause=True)

    images = torch.from_numpy(images)
    t_heatmaps = torch.from_numpy(t_heatmaps)
    l_heatmaps = torch.from_numpy(l_heatmaps)
    b_heatmaps = torch.from_numpy(b_heatmaps)
    r_heatmaps = torch.from_numpy(r_heatmaps)
    ct_heatmaps = torch.from_numpy(ct_heatmaps)
    t_regrs = torch.from_numpy(t_regrs)
    l_regrs = torch.from_numpy(l_regrs)
    b_regrs = torch.from_numpy(b_regrs)
    r_regrs = torch.from_numpy(r_regrs)
    t_tags = torch.from_numpy(t_tags)
    l_tags = torch.from_numpy(l_tags)
    b_tags = torch.from_numpy(b_tags)
    r_tags = torch.from_numpy(r_tags)
    ct_tags = torch.from_numpy(ct_tags)
    tag_masks = torch.from_numpy(tag_masks)

    return {
        "xs": [images, t_tags, l_tags, b_tags, r_tags, ct_tags],
        "ys": [
            t_heatmaps, l_heatmaps, b_heatmaps, r_heatmaps, ct_heatmaps,
            tag_masks, t_regrs, l_regrs, b_regrs, r_regrs
        ]
    }, k_ind
Beispiel #10
0
    def debug(self, batch, output, iter_id):
        opt = self.opt
        # reg = output['reg'] if opt.reg_offset else None
        reg = output['reg'][0:1] 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 = ctdet_decode(output['hm'][0:1],
                            output['wh'][0:1],
                            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

        # FIXME: change from tensor to list and then reshape
        # dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])

        # batch['meta_gt_det'] = [128, 128, 6]
        gt_det = batch['meta_gt_det'][0:1]
        gt_det = np.array(gt_det, dtype=np.float32) if len(gt_det) > 0 else \
            np.zeros((1, 6), dtype=np.float32)
        dets_gt = gt_det.reshape(1, -1, dets.shape[2])
        # print(batch['meta_img_id'][0:1])

        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))
            elif opt.debug == 5:
                debugger.show_all_imgs(pause=opt.pause,
                                       logger=self.logger,
                                       step=iter_id)
            else:
                debugger.show_all_imgs(pause=opt.pause, step=iter_id)
Beispiel #11
0
    def debug(self, batch, output, iter_id):
        opt = self.opt
        reg = output['reg'] if opt.reg_offset else None
        dets = rodet_decode(output['hm'],
                            output['wh'],
                            output['angle'],
                            reg=reg,
                            cat_spec_wh=opt.cat_spec_wh,
                            cat_spec_angle=opt.cat_spec_angle,
                            K=opt.K)
        dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
        dets[:, :, :4] *= opt.down_ratio

        # dets_gt_dense = rodet_decode(
        #   batch['hm'], batch['dense_wh'], batch['dense_angle'], reg=reg,
        #   cat_spec_wh=opt.cat_spec_wh, cat_spec_angle=opt.cat_spec_angle, K=opt.K)
        # dets_gt_dense = dets_gt_dense.detach().cpu().numpy().reshape(1, -1, dets_gt_dense.shape[2])
        # dets_gt_dense[:, :, :4] *= opt.down_ratio

        dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])
        img_name = batch['meta']['img_name']
        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_ang_mask = debugger.gen_colormap(batch['dense_angle_mask'][i].detach().cpu().numpy())
            # gt_ang = debugger.gen_colormap(batch['dense_angle'][i].detach().cpu().numpy())
            gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
            debugger.add_blend_img(img, pred, '{}_pred_hm'.format(img_name))
            # debugger.add_blend_img(img, gt_ang_mask, 'gt_angle_mask')
            # debugger.add_blend_img(img, gt_ang, 'gt_angle')
            debugger.add_blend_img(img, gt, '{}_gt_hm'.format(img_name))
            debugger.add_img(img, img_id='{}_out_pred'.format(img_name))
            for k in range(len(dets[i])):
                if dets[i, k, 5] > opt.center_thresh:
                    # print("pred dets add_rbbox=======================")
                    debugger.add_rbbox(dets[i, k, :5],
                                       dets[i, k, -1],
                                       dets[i, k, 5],
                                       show_txt=False,
                                       img_id='{}_out_pred'.format(img_name))

            # debugger.add_img(img, img_id='{}_dets_gt_dense'.format(img_name))
            # for k in range(len(dets_gt_dense[i])):
            #   if dets_gt_dense[i, k, 5] > opt.center_thresh:
            #     # print("pred dets add_rbbox=======================")
            #     debugger.add_rbbox(dets_gt_dense[i, k, :5], dets_gt_dense[i, k, -1],
            #                        dets_gt_dense[i, k, 5], show_txt=False, img_id='{}_dets_gt_dense'.format(img_name))

            debugger.add_img(img, img_id='{}_out_gt'.format(img_name))
            for k in range(len(dets_gt[i])):
                if dets_gt[i, k, 5] > opt.center_thresh:
                    # print("GT add_rbbox=======================")
                    # 说明add_rbbox(self, rbbox, cat, conf=1, show_txt=True, img_id='default')
                    # gt格式 gt_det.append([ct[0], ct[1], w, h, a, 1, cls_id])
                    debugger.add_rbbox(dets_gt[i, k, :5],
                                       dets_gt[i, k, -1],
                                       dets_gt[i, k, 5],
                                       show_txt=False,
                                       img_id='{}_out_gt'.format(img_name))

            if opt.debug == 4:
                debugger.save_all_imgs(opt.debug_dir,
                                       prefix='{}'.format(iter_id))
            else:
                debugger.show_all_imgs(pause=True)
Beispiel #12
0
    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)