lineless_table_process.py

# ------------------------------------------------------------------------------
# Part of implementation is adopted from CenterNet,
# made publicly available under the MIT License at https://github.com/xingyizhou/CenterNet.git
# ------------------------------------------------------------------------------

import cv2
import numpy as np
import shapely
import torch
import torch.nn as nn
from shapely.geometry import MultiPoint, Point, Polygon


def _gather_feat(feat, ind, mask=None):
    # mandatory
    dim = feat.size(2)
    ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
    feat = feat.gather(1, ind)
    if mask is not None:
        mask = mask.unsqueeze(2).expand_as(feat)
        feat = feat[mask]
        feat = feat.view(-1, dim)
    return feat


def _tranpose_and_gather_feat(feat, ind):
    # mandatory
    feat = feat.permute(0, 2, 3, 1).contiguous()
    feat = feat.view(feat.size(0), -1, feat.size(3))
    feat = _gather_feat(feat, ind)
    return feat


def _get_4ps_feat(cc_match, output):
    # mandatory
    if isinstance(output, dict):
        feat = output['cr']
    else:
        feat = output
    feat = feat.permute(0, 2, 3, 1).contiguous()
    feat = feat.contiguous().view(feat.size(0), -1, feat.size(3))
    feat = feat.unsqueeze(3).expand(
        feat.size(0), feat.size(1), feat.size(2), 4)

    dim = feat.size(2)

    cc_match = cc_match.unsqueeze(2).expand(
        cc_match.size(0), cc_match.size(1), dim, cc_match.size(2))
    if not (isinstance(output, dict)):
        cc_match = torch.where(
            cc_match < feat.shape[1], cc_match, (feat.shape[0] - 1)
            * torch.ones(cc_match.shape).to(torch.int64).cuda())
        cc_match = torch.where(
            cc_match >= 0, cc_match,
            torch.zeros(cc_match.shape).to(torch.int64).cuda())
    feat = feat.gather(1, cc_match)
    return feat


def _nms(heat, name, kernel=3):
    pad = (kernel - 1) // 2

    hmax = nn.functional.max_pool2d(
        heat, (kernel, kernel), stride=1, padding=pad)
    # save_map(hmax.cpu().numpy()[0],name)
    keep = (hmax == heat).float()
    return heat * keep, keep


def _topk(scores, K=40):
    batch, cat, height, width = scores.size()

    topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)

    topk_inds = topk_inds % (
        torch.Tensor([height]).to(torch.int64).cuda()
        * torch.Tensor([width]).to(torch.int64).cuda())
    topk_ys = (topk_inds / torch.Tensor([width]).cuda()).int().float()
    topk_xs = (topk_inds
               % torch.Tensor([width]).to(torch.int64).cuda()).int().float()

    topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
    topk_clses = (topk_ind // K).int()
    topk_inds = _gather_feat(topk_inds.view(batch, -1, 1),
                             topk_ind).view(batch, K)
    topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)

    return topk_score, topk_inds, topk_clses, topk_ys, topk_xs


def corner_decode(mk, st_reg, mk_reg=None, K=400):
    batch, cat, height, width = mk.size()
    mk, keep = _nms(mk, 'mk.0.maxpool')
    scores, inds, clses, ys, xs = _topk(mk, K=K)
    if mk_reg is not None:
        reg = _tranpose_and_gather_feat(mk_reg, inds)
        reg = reg.view(batch, K, 2)
        xs = xs.view(batch, K, 1) + reg[:, :, 0:1]
        ys = ys.view(batch, K, 1) + reg[:, :, 1:2]
    else:
        xs = xs.view(batch, K, 1) + 0.5
        ys = ys.view(batch, K, 1) + 0.5
    scores = scores.view(batch, K, 1)
    st_Reg = _tranpose_and_gather_feat(st_reg, inds)
    bboxes_vec = [
        xs - st_Reg[..., 0:1], ys - st_Reg[..., 1:2], xs - st_Reg[..., 2:3],
        ys - st_Reg[..., 3:4], xs - st_Reg[..., 4:5], ys - st_Reg[..., 5:6],
        xs - st_Reg[..., 6:7], ys - st_Reg[..., 7:8]
    ]
    bboxes = torch.cat(bboxes_vec, dim=2)
    corner_dict = {
        'scores': scores,
        'inds': inds,
        'ys': ys,
        'xs': xs,
        'gboxes': bboxes
    }
    return scores, inds, ys, xs, bboxes, corner_dict


def ctdet_4ps_decode(heat,
                     wh,
                     ax,
                     cr,
                     corner_dict=None,
                     reg=None,
                     cat_spec_wh=False,
                     K=100,
                     wiz_rev=False):

    batch, cat, height, width = heat.size()
    # heat = torch.sigmoid(heat)
    # perform nms on heatmaps
    heat, keep = _nms(heat, 'hm.0.maxpool')

    scores, inds, clses, ys, xs = _topk(heat, K=K)
    if reg is not None:
        reg = _tranpose_and_gather_feat(reg, inds)
        reg = reg.view(batch, K, 2)
        xs = xs.view(batch, K, 1) + reg[:, :, 0:1]
        ys = ys.view(batch, K, 1) + reg[:, :, 1:2]
    else:
        xs = xs.view(batch, K, 1) + 0.5
        ys = ys.view(batch, K, 1) + 0.5
    wh = _tranpose_and_gather_feat(wh, inds)
    ax = _tranpose_and_gather_feat(ax, inds)

    if cat_spec_wh:
        wh = wh.view(batch, K, cat, 8)
        clses_ind = clses.view(batch, K, 1, 1).expand(batch, K, 1, 8).long()
        wh = wh.gather(2, clses_ind).view(batch, K, 8)
    else:
        wh = wh.view(batch, K, 8)
    clses = clses.view(batch, K, 1).float()
    scores = scores.view(batch, K, 1)

    bboxes_vec = [
        xs - wh[..., 0:1], ys - wh[..., 1:2], xs - wh[..., 2:3],
        ys - wh[..., 3:4], xs - wh[..., 4:5], ys - wh[..., 5:6],
        xs - wh[..., 6:7], ys - wh[..., 7:8]
    ]
    bboxes = torch.cat(bboxes_vec, dim=2)

    cc_match = torch.cat(
        [(xs - wh[..., 0:1]) + width * torch.round(ys - wh[..., 1:2]),
         (xs - wh[..., 2:3]) + width * torch.round(ys - wh[..., 3:4]),
         (xs - wh[..., 4:5]) + width * torch.round(ys - wh[..., 5:6]),
         (xs - wh[..., 6:7]) + width * torch.round(ys - wh[..., 7:8])],
        dim=2)

    cc_match = torch.round(cc_match).to(torch.int64)

    cr_feat = _get_4ps_feat(cc_match, cr)
    cr_feat = cr_feat.sum(axis=3)

    detections = torch.cat([bboxes, scores, clses], dim=2)

    return detections, keep, ax, cr_feat


def get_3rd_point(a, b):
    direct = a - b
    return b + np.array([-direct[1], direct[0]], dtype=np.float32)


def affine_transform(pt, t):
    new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32).T
    new_pt = np.dot(t, new_pt)
    return new_pt[:2]


def get_dir(src_point, rot_rad):
    sn, cs = np.sin(rot_rad), np.cos(rot_rad)

    src_result = [0, 0]
    src_result[0] = src_point[0] * cs - src_point[1] * sn
    src_result[1] = src_point[0] * sn + src_point[1] * cs

    return src_result


def get_affine_transform(center,
                         scale,
                         rot,
                         output_size,
                         shift=np.array([0, 0], dtype=np.float32),
                         inv=0):
    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
        scale = np.array([scale, scale], dtype=np.float32)

    scale_tmp = scale
    src_w = scale_tmp[0]
    dst_w = output_size[0]
    dst_h = output_size[1]

    rot_rad = np.pi * rot / 180
    src_dir = get_dir([0, src_w * -0.5], rot_rad)
    dst_dir = np.array([0, dst_w * -0.5], np.float32)

    src = np.zeros((3, 2), dtype=np.float32)
    dst = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center + scale_tmp * shift  # [0,0] #
    src[1, :] = center + src_dir + scale_tmp * shift  # scale #
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]  # [0,0] #
    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5],
                         np.float32) + dst_dir  # output_size #

    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])

    if inv:
        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
    else:
        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))

    return trans


def get_affine_transform_upper_left(center,
                                    scale,
                                    rot,
                                    output_size,
                                    shift=np.array([0, 0], dtype=np.float32),
                                    inv=0):
    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
        scale = np.array([scale, scale], dtype=np.float32)

    src = np.zeros((3, 2), dtype=np.float32)
    dst = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center
    dst[0, :] = [0, 0]
    if center[0] < center[1]:
        src[1, :] = [scale[0], center[1]]
        dst[1, :] = [output_size[0], 0]
    else:
        src[1, :] = [center[0], scale[0]]
        dst[1, :] = [0, output_size[0]]
    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])

    if inv:
        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
    else:
        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))

    return trans


def transform_preds(coords, center, scale, output_size, rot=0):
    target_coords = np.zeros(coords.shape)
    trans = get_affine_transform(center, scale, rot, output_size, inv=1)
    for p in range(coords.shape[0]):
        target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
    return target_coords


def transform_preds_upper_left(coords, center, scale, output_size, rot=0):
    target_coords = np.zeros(coords.shape)

    trans = get_affine_transform_upper_left(
        center, scale, rot, output_size, inv=1)
    for p in range(coords.shape[0]):
        target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
    return target_coords


def ctdet_4ps_post_process_upper_left(dets, c, s, h, w, num_classes, rot=0):
    # dets: batch x max_dets x dim
    # return 1-based class det dict
    ret = []
    for i in range(dets.shape[0]):
        top_preds = {}
        dets[i, :, 0:2] = transform_preds_upper_left(dets[i, :, 0:2], c[i],
                                                     s[i], (w, h), rot)
        dets[i, :, 2:4] = transform_preds_upper_left(dets[i, :, 2:4], c[i],
                                                     s[i], (w, h), rot)
        dets[i, :, 4:6] = transform_preds_upper_left(dets[i, :, 4:6], c[i],
                                                     s[i], (w, h), rot)
        dets[i, :, 6:8] = transform_preds_upper_left(dets[i, :, 6:8], c[i],
                                                     s[i], (w, h), rot)
        classes = dets[i, :, -1]
        for j in range(num_classes):
            inds = (classes == j)
            tmp_top_pred = [
                dets[i, inds, :8].astype(np.float32),
                dets[i, inds, 8:9].astype(np.float32)
            ]
            top_preds[j + 1] = np.concatenate(tmp_top_pred, axis=1).tolist()
        ret.append(top_preds)
    return ret


def ctdet_corner_post_process(corner_st_reg, c, s, h, w, num_classes):
    for i in range(corner_st_reg.shape[0]):
        corner_st_reg[i, :, 0:2] = transform_preds(corner_st_reg[i, :, 0:2],
                                                   c[i], s[i], (w, h))
        corner_st_reg[i, :, 2:4] = transform_preds(corner_st_reg[i, :, 2:4],
                                                   c[i], s[i], (w, h))
        corner_st_reg[i, :, 4:6] = transform_preds(corner_st_reg[i, :, 4:6],
                                                   c[i], s[i], (w, h))
        corner_st_reg[i, :, 6:8] = transform_preds(corner_st_reg[i, :, 6:8],
                                                   c[i], s[i], (w, h))
        corner_st_reg[i, :, 8:10] = transform_preds(corner_st_reg[i, :, 8:10],
                                                    c[i], s[i], (w, h))
    return corner_st_reg


def merge_outputs(detections):
    # thresh_conf, thresh_min, thresh_max = 0.1, 0.5, 0.7
    num_classes, max_per_image = 2, 3000
    results = {}
    for j in range(1, num_classes + 1):
        results[j] = np.concatenate([detection[j] for detection in detections],
                                    axis=0).astype(np.float32)
    scores = np.hstack([results[j][:, 8] for j in range(1, num_classes + 1)])
    if len(scores) > max_per_image:
        kth = len(scores) - max_per_image
        thresh = np.partition(scores, kth)[kth]
        for j in range(1, num_classes + 1):
            keep_inds = (results[j][:, 8] >= thresh)
            results[j] = results[j][keep_inds]
    return results


def filter(results, logi, ps):
    # this function select boxes
    batch_size, feat_dim = logi.shape[0], logi.shape[2]

    num_valid = sum(results[1][:, 8] >= 0.15)

    slct_logi = np.zeros((batch_size, num_valid, feat_dim), dtype=np.float32)
    slct_dets = np.zeros((batch_size, num_valid, 8), dtype=np.int32)
    for i in range(batch_size):
        for j in range(num_valid):
            slct_logi[i, j, :] = logi[i, j, :].cpu()
            slct_dets[i, j, :] = ps[i, j, :].cpu()

    return torch.Tensor(slct_logi).cuda(), torch.Tensor(slct_dets).cuda()


def normalized_ps(ps, vocab_size):
    ps = torch.round(ps).to(torch.int64)
    ps = torch.where(ps < vocab_size, ps, (vocab_size - 1)
                     * torch.ones(ps.shape).to(torch.int64).cuda())
    ps = torch.where(ps >= 0, ps, torch.zeros(ps.shape).to(torch.int64).cuda())
    return ps


def process_detect_output(output, meta):
    K, MK = 3000, 5000
    num_classes = 2
    scale = 1.0

    hm = output['hm'].sigmoid_()
    wh = output['wh']
    reg = output['reg']
    st = output['st']
    ax = output['ax']
    cr = output['cr']

    scores, inds, ys, xs, st_reg, corner_dict = corner_decode(
        hm[:, 1:2, :, :], st, reg, K=MK)
    dets, keep, logi, cr = ctdet_4ps_decode(
        hm[:, 0:1, :, :], wh, ax, cr, corner_dict, reg=reg, K=K, wiz_rev=False)
    raw_dets = dets
    dets = dets.detach().cpu().numpy()
    dets = dets.reshape(1, -1, dets.shape[2])
    dets = ctdet_4ps_post_process_upper_left(dets.copy(),
                                             [meta['c'].cpu().numpy()],
                                             [meta['s']], meta['out_height'],
                                             meta['out_width'], 2)
    for j in range(1, num_classes + 1):
        dets[0][j] = np.array(dets[0][j], dtype=np.float32).reshape(-1, 9)
        dets[0][j][:, :8] /= scale
    dets = dets[0]
    detections = [dets]

    logi = logi + cr
    results = merge_outputs(detections)
    slct_logi_feat, slct_dets_feat = filter(results, logi, raw_dets[:, :, :8])
    slct_dets_feat = normalized_ps(slct_dets_feat, 256)
    slct_output_dets = results[1][:slct_logi_feat.shape[1], :8]

    return slct_logi_feat, slct_dets_feat, slct_output_dets


def process_logic_output(logi):
    logi_floor = logi.floor()
    dev = logi - logi_floor
    logi = torch.where(dev > 0.5, logi_floor + 1, logi_floor)

    return logi


def load_lore_model(model, checkpoint, mtype):
    state_dict_ = checkpoint['state_dict']
    state_dict = {}
    # convert data_parallal to model
    for k in state_dict_:
        if k.startswith('module') and not k.startswith('module_list'):
            state_dict[k[7:]] = state_dict_[k]
        else:
            if mtype == 'model':
                if k.startswith('model'):
                    state_dict[k[6:]] = state_dict_[k]
                else:
                    continue
            else:
                if k.startswith('processor'):
                    state_dict[k[10:]] = state_dict_[k]
                else:
                    continue
    model_state_dict = model.state_dict()
    # check loaded parameters and created model parameters
    for k in state_dict:
        if k in model_state_dict:
            if state_dict[k].shape != model_state_dict[k].shape:
                print('Skip loading parameter {}, required shape{}, '
                      'loaded shape{}.'.format(k, model_state_dict[k].shape,
                                               state_dict[k].shape))
                state_dict[k] = model_state_dict[k]
        else:
            print('Drop parameter {}.'.format(k))
    for k in model_state_dict:
        if not (k in state_dict):
            print('No param {}.'.format(k))
            state_dict[k] = model_state_dict[k]
    model.load_state_dict(state_dict, strict=False)