undistort_detection_result.py

undistort_detection_result.py

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
#
# Demo of un-distorting 2D bbox detection result instead of image for potential speedup
#    114       484      67472.0    139.4      0.2          dets_tweak = tweak_dets_by_undistort_points_vec(dets, K=K_cam, distortion=distortion, R=undistort_R, P=undistort_P)
#    117       484    2853099.0   5894.8      6.3          im_undistort = distort_fun(im)
#
# Ref:
# https://stackoverflow.com/questions/8499984/how-to-undistort-points-in-camera-shot-coordinates-and-obtain-corresponding-undi
# FIXME, tried to use mapping array, but failed


import os
import cv2
import time
import numpy as np
import json_tricks as json

from det.inference import DetInfer
from det.cfgs.config import det_cfgs

from depthwheel.cameras.undistortion import StereoImageUndistort, MonoImageUndistort


def tweak_dets_by_mapping(dets, map1, map2, height, width):
    dets = dets.copy()
    for i, det in enumerate(dets):
        x1, y1, x2, y2 = [int(v) for v in det[:4]]
        x1 = max(x1, 0)
        y1 = max(y1, 0)
        x2 = max(x2, 0)
        y2 = max(y2, 0)

        x1 = min(x1, width-1)
        x2 = min(x2, width-1)
        y1 = min(y1, height-1)
        y2 = min(y2, height-1)

        x1_ = map1[y1, x1]
        y1_ = map2[y1, x1]
        x2_ = map1[y2, x2]
        y2_ = map2[y2, x2]

        dets[i][:4] = x1_, y1_, x2_, y2_
    return dets


def tweak_dets_by_undistort_points(dets, K, distortion, R=None, P=None):
    dets = dets.copy()
    for i, det in enumerate(dets):
        x1, y1, x2, y2 = det[:4]
        x1_, y1_ = undistort_uv((x1, y1), K, distortion, R, P)
        x2_, y2_ = undistort_uv((x2, y2), K, distortion, R, P)
        dets[i][:4] = x1_, y1_, x2_, y2_
    return dets


def tweak_dets_by_undistort_points_vec(dets, K, distortion, R=None, P=None):
    dets = dets.copy()
    uv_s = dets[:, :4].reshape((-1, 2))
    uv_s = undistort_uv(uv_s, K, distortion, R, P).reshape((-1, 4))
    for i in range(len(dets)):
        dets[i][:4] = uv_s[i].ravel()
    return dets


def undistort_uv(uv, K, distort, R, P):
    if P is None:
        P = K
    uv = cv2.undistortPoints(uv, K, distort, R=R, P=P)
    return uv.ravel()


@profile
def run():
    imgs_dir = '/tmp/2019-10-23-14-17-32/dev1_cam2/'

    calib_f = '../../cfgs/dm_vv6_calib_distort_ningsheng.json'

    with open(calib_f, 'r') as f:
        calib = json.load(f)

    K_cam = calib['K_cam_left_distort']
    distortion = calib['distortion_left']

    # # MONO
    # undistort = MonoImageUndistort.from_mat(K_cam=K_cam, K_cam_new=K_cam,
    #                                         distortion=distortion, resolution=calib['resolution'])
    # map1, map2 = undistort.map_x, undistort.map_y
    # def distort_fun(im):
    #     return undistort.process_image(im)
    # undistort_R = None
    # undistort_P = K_cam

    # STEREO
    stereo_undistort = StereoImageUndistort(distort_f=calib_f)
    map1, map2 = stereo_undistort.map_d['left']
    def distort_fun(im):
        return stereo_undistort.process_image(im, is_left=True)
    undistort_R = stereo_undistort.R1_to_rectify
    undistort_P = stereo_undistort.P1

    imgs_fn_list = [fn for fn in sorted(os.listdir(imgs_dir))]

    det_infer = DetInfer(det_cfgs)

    for idx, fn in enumerate(imgs_fn_list):
        im = cv2.imread(os.path.join(imgs_dir, fn))
        tic = time.time()
        dets, _ = det_infer.inference(im)
        dets = np.array(dets)
        print("infer time:", time.time()-tic)

        im_s_direct = det_infer.vis_result(im.copy(), dets, thickness=3)

        dets_tweak_mapping = tweak_dets_by_mapping(dets, map1, map2, height=720, width=1280)
        dets_tweak_ = tweak_dets_by_undistort_points(dets, K=K_cam, distortion=distortion, R=undistort_R, P=undistort_P)
        dets_tweak = tweak_dets_by_undistort_points_vec(dets, K=K_cam, distortion=distortion, R=undistort_R, P=undistort_P)
        assert np.allclose(dets_tweak, dets_tweak_)

        im_undistort = distort_fun(im)
        im_s_undistort = det_infer.vis_result(im_undistort, dets_tweak, thickness=3)

        im_s = np.vstack((im_s_direct, im_s_undistort))
        im_s = cv2.resize(im_s, dsize=None, fx=0.4, fy=0.4)

        cv2.imshow("results", im_s)
        cv2.waitKey(1)


if __name__ == '__main__':
    run()

Note, 2D detection model, calib file, image folder, some helper functions, and possible speedup are all NOT provided here, as only the tweak_dets_by_undistort_points_vec and distort_fun matter.