correspondence_database.py

import random
import time
import numpy as np
import os
import cv2
from skimage.io import imread, imsave

from dataset.GL3D_dataset import get_gl3d_dataset
from utils.base_utils import perspective_transform, read_pickle, save_pickle


def worker_init_fn(worker_id):
    np.random.seed(worker_id + (int(round(time.time() * 1000) % (2 ** 16))))


def scale_transform_img(img, min_ratio=0.15, max_ratio=0.25, base_ratio=2, flip=True):
    h, w = img.shape[0], img.shape[1]
    pts0 = np.asarray([[0, 0], [w, 0], [w, h], [0, h]], np.float32)

    scale_ratio = base_ratio ** (-np.random.uniform(min_ratio, max_ratio))
    if np.random.random() < 0.5 and flip:
        scale_ratio = 1.0 / scale_ratio
    center = np.mean(pts0, 0, keepdims=True)
    pts1 = (pts0 - center) * scale_ratio + center
    if scale_ratio > 1:
        min_pt = np.min(pts1, 0)  # <0
        max_pt = np.max(pts1, 0)  # >w,h
        min_w, min_h = -(max_pt - np.asarray([w, h]))
        max_w, max_h = -min_pt
    else:
        min_pt = np.min(pts1, 0)  # >0
        max_pt = np.max(pts1, 0)  # <w,h
        min_w, min_h = -min_pt
        max_w, max_h = np.asarray([w, h]) - max_pt

    offset_h = np.random.uniform(min_h, max_h)
    offset_w = np.random.uniform(min_w, max_w)
    pts1 += np.asarray([[offset_w, offset_h]], np.float32)

    th, tw = h, w  # int(h * scale_ratio), int(w * scale_ratio)
    H = cv2.getPerspectiveTransform(pts0.astype(np.float32), pts1.astype(np.float32))

    img1 = cv2.warpPerspective(img, H, (tw, th), flags=cv2.INTER_LINEAR)
    return img1, H


def perspective_transform_img(img, perspective_type='lr', min_ratio=0.05, max_ratio=0.1):
    h, w = img.shape[0], img.shape[1]
    pts0 = np.asarray([[0, 0], [w, 0], [w, h], [0, h]], np.float32)
    pts1 = np.asarray([[0, 0], [w, 0], [w, h], [0, h]], np.float32)

    # left right
    if perspective_type == 'lr':
        val = h * np.random.uniform(min_ratio, max_ratio)
        if np.random.random() < 0.5: val *= -1
        pts1[0, 1] -= val
        pts1[1, 1] += val
        pts1[2, 1] -= val
        pts1[3, 1] += val

        val = h * np.random.uniform(min_ratio, max_ratio)
        pts1[0, 0] += val
        pts1[1, 0] -= val
        pts1[2, 0] -= val
        pts1[3, 0] += val
    else:  # 'ud'
        val = w * np.random.uniform(min_ratio, max_ratio)
        if np.random.random() < 0.5: val *= -1
        pts1[0, 0] += val
        pts1[1, 0] -= val
        pts1[2, 0] += val
        pts1[3, 0] -= val

        val = h * np.random.uniform(min_ratio, max_ratio)
        pts1[0, 1] += val
        pts1[1, 1] += val
        pts1[2, 1] -= val
        pts1[3, 1] -= val

    pts1 = pts1 - np.min(pts1, 0, keepdims=True)
    tw, th = np.max(pts1, 0)
    H = cv2.getPerspectiveTransform(pts0.astype(np.float32), pts1.astype(np.float32))
    img1 = cv2.warpPerspective(img, H, (tw, th), flags=cv2.INTER_LINEAR)
    return img1, H


def rotate_transform_img(img, min_angle=0, max_angle=360, random_flip=False):
    h, w = img.shape[0], img.shape[1]

    pts0 = np.asarray([[0, 0], [w, 0], [w, h], [0, h]], np.float32)
    center = np.mean(pts0, 0, keepdims=True)
    theta = np.random.uniform(min_angle / 180 * np.pi, max_angle / 180 * np.pi)
    if random_flip and np.random.random() < 0.5: theta = -theta
    R = np.asarray([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]], np.float32)
    pts1 = (pts0 - center) @ R + center
    H = cv2.getPerspectiveTransform(pts0.astype(np.float32), pts1.astype(np.float32))

    img1 = cv2.warpPerspective(img, H, (w, h))
    return img1, H


class CorrespondenceDatabase:
    @staticmethod
    def get_SUN2012_image_paths():
        img_dir = os.path.join('data', 'SUN2012Images', 'JPEGImages')
        img_pths = [os.path.join(img_dir, fn) for fn in os.listdir(img_dir)]
        return img_pths

    @staticmethod
    def get_COCO_image_paths():
        img_dir = os.path.join('data', 'coco', 'train2014')
        img_pths = [os.path.join(img_dir, fn) for fn in os.listdir(img_dir)]
        img_dir = os.path.join('data', 'coco', 'val2014')
        img_pths += [os.path.join(img_dir, fn) for fn in os.listdir(img_dir)]
        return img_pths

    @staticmethod
    def generate_homography_database(img_list):
        return [{'type': 'homography', 'img_pth': img_pth} for img_pth in img_list]

    @staticmethod
    def get_hpatch_sequence_database(name='resize', max_size=480):
        """
        Get hpatches_resize if it exists,
            else generate one
        """

        def resize_and_save(pth_in, max_size, pth_out):
            img = imread(pth_in)
            h, w = img.shape[:2]
            ratio = max_size / max(h, w)
            h, w = int(h * ratio), int(w * ratio)
            img = cv2.resize(img, (w, h))
            imsave(pth_out, img)
            return ratio

        root_dir = os.path.join('data', 'hpatches_sequence')
        output_dir = os.path.join('data', 'hpatches_{}'.format(name))
        pkl_file = os.path.join(output_dir, 'info.pkl')
        if os.path.exists(pkl_file):
            return read_pickle(pkl_file)

        if not os.path.exists(output_dir): os.mkdir(output_dir)
        illumination_dataset = []
        viewpoint_dataset = []
        for dir in os.listdir(root_dir):
            if not os.path.exists(os.path.join(output_dir, dir)):
                os.mkdir(os.path.join(output_dir, dir))

            img_pattern = os.path.join(root_dir, dir, '{}.ppm')
            hmg_pattern = os.path.join(root_dir, dir, 'H_1_{}')
            omg_pattern = os.path.join(output_dir, dir, '{}.png')

            ratio0 = resize_and_save(img_pattern.format(1), max_size, omg_pattern.format(1))
            # resize image
            for k in range(2, 7):
                ratio1 = resize_and_save(img_pattern.format(k), max_size, omg_pattern.format(k))
                H = np.loadtxt(hmg_pattern.format(k))
                H = np.matmul(np.diag([ratio1, ratio1, 1.0]), np.matmul(H, np.diag([1 / ratio0, 1 / ratio0, 1.0])))
                data = {'type': 'hpatch',
                        'img0_pth': omg_pattern.format(1),
                        'img1_pth': omg_pattern.format(k),
                        'H': H}
                if dir.startswith('v'):
                    viewpoint_dataset.append(data)
                if dir.startswith('i'):
                    illumination_dataset.append(data)

        save_pickle([illumination_dataset, viewpoint_dataset], pkl_file)
        return illumination_dataset, viewpoint_dataset

    @staticmethod
    def add_homography_background(img, H):
        img_dir = os.path.join('data', 'SUN2012Images', 'JPEGImages')
        background_pths = [os.path.join(img_dir, fn) for fn in os.listdir(img_dir)]
        bpth = background_pths[np.random.randint(0, len(background_pths))]

        h, w, _ = img.shape
        bimg = cv2.resize(imread(bpth), (w, h))
        if len(bimg.shape) == 2: bimg = np.repeat(bimg[:, :, None], 3, axis=2)
        if bimg.shape[2] > 3: bimg = bimg[:, :, :3]
        msk_tgt = cv2.warpPerspective(np.ones([h, w], np.uint8), H, (w, h), flags=cv2.INTER_NEAREST).astype(np.bool)
        img[np.logical_not(msk_tgt)] = bimg[np.logical_not(msk_tgt)]
        return img

    @staticmethod
    def make_hpatch_transform_database_combine(in_dataset, output_name, transform, add_background, identity=False):
        hpatch_transform_root_dir = os.path.join('data', 'hpatches_{}'.format(output_name))
        if not os.path.exists(hpatch_transform_root_dir): os.mkdir(hpatch_transform_root_dir)

        hpatch_transform_pkl = os.path.join(hpatch_transform_root_dir, 'info.pkl')
        if os.path.exists(hpatch_transform_pkl):
            return read_pickle(hpatch_transform_pkl)

        print('begin making {} dataset'.format(output_name))
        dataset = []
        random.shuffle(in_dataset)
        for in_data in in_dataset:
            pth0 = in_data['img0_pth']
            pth1 = in_data['img1_pth']
            in_dir = pth1.split('/')[-2]
            in_id = pth1.split('/')[-1].split('.')[-2]
            output_dir = os.path.join(hpatch_transform_root_dir, in_dir)
            if not os.path.exists(output_dir): os.mkdir(output_dir)

            img1 = imread(pth1)
            img1, H = transform(img1)

            if add_background: img1 = CorrespondenceDatabase.add_homography_background(img1, H)
            img1_pth = os.path.join(output_dir, '{}.png'.format(in_id))
            imsave(img1_pth, img1)
            if identity:
                data = {'type': 'hpatch',
                        'img0_pth': pth1,
                        'img1_pth': img1_pth,
                        'H': H}
            else:
                data = {'type': 'hpatch',
                        'img0_pth': pth0,
                        'img1_pth': img1_pth,
                        'H': H @ in_data['H']}
            dataset.append(data)

        save_pickle(dataset, hpatch_transform_pkl)
        return dataset

    @staticmethod
    def warp_flow(H, norm_flow):
        h, w, _ = norm_flow.shape
        xs, ys = np.meshgrid(np.arange(w), np.arange(h))
        pts = np.concatenate([xs.reshape([-1, 1]), ys.reshape([-1, 1])], 1)

        norm_flow = norm_flow.reshape([h * w, 2])
        nan_mask = np.sum(np.isnan(norm_flow), 1) > 0
        norm_flow[nan_mask] = (0, 0)
        norm_flow *= np.asarray((w, h))

        norm_flow += pts
        outside_mask = np.logical_or(np.logical_or(norm_flow[:, 0] < 0, norm_flow[:, 0] >= w),
                                     np.logical_or(norm_flow[:, 1] < 0, norm_flow[:, 1] >= h))
        norm_flow = perspective_transform(norm_flow, H)
        norm_flow -= pts
        norm_flow /= np.asarray((w, h))
        norm_flow[nan_mask] = np.nan
        norm_flow[outside_mask] = np.nan

        norm_flow = norm_flow.reshape([h, w, 2])
        return norm_flow

    @staticmethod
    def make_sun3d_all_dataset(img_num=499):
        dataset = []
        for k in range(img_num):
            data = {
                'img0_pth': 'data/sun3d_all/img0/{}.png'.format(k),
                'img1_pth': 'data/sun3d_all/img1/{}.png'.format(k),
                'flow_pth': 'data/sun3d_all/flow_01/{}.npy'.format(k),
                'type': 'sun3d_all'
            }
            dataset.append(data)
        return dataset

    @staticmethod
    def make_transformed_sun3d_dataset(dataset_in, dataset_out_name, warp_fn, add_background=True):
        dir_out = os.path.join('data', dataset_out_name)
        if not os.path.exists(dir_out): os.mkdir(dir_out)
        info_pkl = os.path.join(dir_out, 'info.pkl')
        if os.path.exists(info_pkl):
            return read_pickle(info_pkl)

        if not os.path.exists(os.path.join(dir_out, 'img1')):
            os.mkdir(os.path.join(dir_out, 'img1'))
        if not os.path.exists(os.path.join(dir_out, 'flow_01')):
            os.mkdir(os.path.join(dir_out, 'flow_01'))

        print('begin making {} dataset'.format(dataset_out_name))

        dataset_out = []
        for data_in in dataset_in:
            image_id = data_in['img0_pth'].split('/')[-1].split('.')[0]
            img1 = imread(data_in['img1_pth']).astype(np.uint8)
            flow = np.load(data_in['flow_pth']).astype(np.float32).transpose([1, 2, 0])

            img1, H = warp_fn(img1)
            if add_background:
                img1 = CorrespondenceDatabase.add_homography_background(img1, H)

            flow_warp = CorrespondenceDatabase.warp_flow(H, flow)

            data_out = {
                'img0_pth': data_in['img0_pth'],
                'img1_pth': os.path.join(dir_out, 'img1', '{}.png'.format(image_id)),
                'flow_pth': os.path.join(dir_out, 'flow_01', '{}.npy'.format(image_id)),
                'type': 'sun3d_all'
            }
            imsave(data_out['img1_pth'], img1)
            np.save(data_out['flow_pth'], flow_warp.astype(np.float32))
            dataset_out.append(data_out)

        save_pickle(dataset_out, info_pkl)
        return dataset_out

    def __getattr__(self, item):
        if item == 'coco_set':
            self.coco_set = self.generate_homography_database(self.get_COCO_image_paths())
            print('coco_len {}'.format(len(self.coco_set)))
            return self.coco_set
        elif item== 'hi_set' or item=='hv_set':
            self.hi_set, self.hv_set = self.get_hpatch_sequence_database()
            if item=='hi_set': return self.hi_set
            else: return self.hv_set
        elif item.startswith('erotate') or item.startswith('escale'):
            scaling = lambda img: scale_transform_img(img, 1.5, 2.0)
            rotating = lambda img: rotate_transform_img(img, -90, 90, True)
            escale_set = self.make_hpatch_transform_database_combine(
                self.hv_set, 'escale', scaling, True)
            escale_illm_set = self.make_hpatch_transform_database_combine(
                self.hi_set, 'escale_illm', scaling, True)
            self.escale_set=escale_set+escale_illm_set
            erotate_set = self.make_hpatch_transform_database_combine(
                self.hv_set, 'erotate', rotating, True)
            erotate_illm_set = self.make_hpatch_transform_database_combine(
                self.hi_set, 'erotate_illm', rotating, True)
            self.erotate_set=erotate_set+erotate_illm_set
            if item=='erotate_set':
                return self.erotate_set
            else:
                return self.escale_set
        elif item=='gl3d_set':
            self.gl3d_set=get_gl3d_dataset()
            print('gl3d len {}'.format(len(self.gl3d_set)))
            return self.gl3d_set
        else:
            super(CorrespondenceDatabase, self).__getattribute__(item)

    def __init__(self):
        pass