utils.py

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import cv2
from tqdm import tqdm #_notebook as tqdm
from functools import reduce
import os
from sklearn.model_selection import train_test_split
from scipy.optimize import minimize
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
from math import sin, cos
from scipy.optimize import minimize

PATH = './data/'
CENTER_NET = True
camera_matrix = np.array([[2304.5479, 0,  1686.2379],
                          [0, 2305.8757, 1354.9849],
                          [0, 0, 1]], dtype=np.float32)

camera_matrix_inv = np.linalg.inv(camera_matrix)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def coords2str(coords, names=['yaw', 'pitch', 'roll', 'x', 'y', 'z', 'confidence']):
    s = []
    for c in coords:
        for n in names:
            s.append(str(c.get(n, 0)))
    return ' '.join(s)


def str2coords(s, names=['id', 'yaw', 'pitch', 'roll', 'x', 'y', 'z']):
    '''
    Input:
        s: PredictionString (e.g. from train dataframe)
        names: array of what to extract from the string
    Output:
        list of dicts with keys from `names`
    '''
    coords = []
    for l in np.array(s.split()).reshape([-1, 7]):
        coords.append(dict(zip(names, l.astype('float'))))
        if 'id' in coords[-1]:
            coords[-1]['id'] = int(coords[-1]['id'])
    return coords


def _regr_preprocess(regr_dict, flip=False):
    if flip:
        for k in ['x', 'pitch', 'roll']:
            regr_dict[k] = -regr_dict[k]
    for name in ['x', 'y', 'z']:
        regr_dict[name] = regr_dict[name] / 100
    regr_dict['roll'] = rotate(regr_dict['roll'], np.pi)
    regr_dict['pitch_sin'] = sin(regr_dict['pitch'])
    regr_dict['pitch_cos'] = cos(regr_dict['pitch'])
    regr_dict.pop('pitch')
    regr_dict.pop('id')
    return regr_dict


def imread(path, fast_mode=False):
    img = cv2.imread(path)
    if not fast_mode and img is not None and len(img.shape) == 3:
        img = np.array(img[:, :, ::-1])
    return img


def _regr_back(regr_dict):
    for name in ['x', 'y', 'z']:
        regr_dict[name] = regr_dict[name] * 100
    regr_dict['roll'] = rotate(regr_dict['roll'], -np.pi)

    pitch_sin = regr_dict['pitch_sin'] / np.sqrt(regr_dict['pitch_sin'] ** 2 + regr_dict['pitch_cos'] ** 2)
    pitch_cos = regr_dict['pitch_cos'] / np.sqrt(regr_dict['pitch_sin'] ** 2 + regr_dict['pitch_cos'] ** 2)
    regr_dict['pitch'] = np.arccos(pitch_cos) * np.sign(pitch_sin)
    return regr_dict


def preprocess_mask(img):
    img = img[img.shape[0] // 2:]
    bg = np.ones_like(img) * img.mean(1, keepdims=True).astype(img.dtype)
    bg = bg[:, :img.shape[1] // 6]
    img = np.concatenate([bg, img, bg], 1)
    img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))
    return img


def preprocess_image(img, flip=False):
    img = img[img.shape[0] // 2:]
    bg = np.ones_like(img) * img.mean(1, keepdims=True).astype(img.dtype)
    bg = bg[:, :img.shape[1] // 6]
    img = np.concatenate([bg, img, bg], 1)
    img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))
    if flip:
        img = img[:, ::-1]
    return (img / 255).astype('float32')


def rotate(x, angle):
    x = x + angle
    x = x - (x + np.pi) // (2 * np.pi) * 2 * np.pi
    return x


IMG_WIDTH = 2048 - 512
IMG_HEIGHT = IMG_WIDTH // 16 * 5
MODEL_SCALE = 4


def get_img_coords(s):
    '''
    Input is a PredictionString (e.g. from train dataframe)
    Output is two arrays:
        xs: x coordinates in the image
        ys: y coordinates in the image
    '''
    coords = str2coords(s)
    xs = [c['x'] for c in coords]
    ys = [c['y'] for c in coords]
    zs = [c['z'] for c in coords]
    P = np.array(list(zip(xs, ys, zs))).T
    img_p = np.dot(camera_matrix, P).T
    img_p[:, 0] /= img_p[:, 2]
    img_p[:, 1] /= img_p[:, 2]
    img_xs = img_p[:, 0]
    img_ys = img_p[:, 1]
    img_zs = img_p[:, 2] # z = Distance from the camera
    return img_xs, img_ys


def get_mask_and_regr(img, labels, sigma=1, flip=False):
    mask = np.zeros([IMG_HEIGHT // MODEL_SCALE, IMG_WIDTH // MODEL_SCALE], dtype='float32')
    regr_names = ['x', 'y', 'z', 'yaw', 'pitch', 'roll']
    regr = np.zeros([IMG_HEIGHT // MODEL_SCALE, IMG_WIDTH // MODEL_SCALE, 7], dtype='float32')
    coords = str2coords(labels)
    xs, ys = get_img_coords(labels)
    count = 0
    heatmap = np.zeros((IMG_HEIGHT // MODEL_SCALE, IMG_WIDTH // MODEL_SCALE))
    for x, y, regr_dict in zip(xs, ys, coords):
        x, y = y, x
        x = (x - img.shape[0] // 2) * IMG_HEIGHT / (img.shape[0] // 2) / MODEL_SCALE
        x = np.round(x).astype('int')
        y = (y + img.shape[1] // 6) * IMG_WIDTH / (img.shape[1] * 4 / 3) / MODEL_SCALE
        y = np.round(y).astype('int')

        if x >= 0 and x < IMG_HEIGHT // MODEL_SCALE and y >= 0 and y < IMG_WIDTH // MODEL_SCALE:
            mask[x, y] = 1
            mask_try, heatmap_i = gaussian_kernel(IMG_HEIGHT // MODEL_SCALE, IMG_WIDTH // MODEL_SCALE, x, y, sigma)
            assert mask_try[x, y] == 1
            assert np.abs(heatmap_i[x, y] - 1) < 1e-6
            heatmap = np.maximum(heatmap, heatmap_i)
            regr_dict = _regr_preprocess(regr_dict, flip)
            regr[x, y] = [regr_dict[n] for n in sorted(regr_dict)]
            count += 1
    if flip:
        mask = np.array(mask[:, ::-1])
        regr = np.array(regr[:, ::-1])
        heatmap = np.array((heatmap[:, ::-1]))

    return mask, regr, heatmap


def get_mesh(batch_size, shape_x, shape_y):
    mg_x, mg_y = np.meshgrid(np.linspace(0, 1, shape_y), np.linspace(0, 1, shape_x))
    mg_x = np.tile(mg_x[None, None, :, :], [batch_size, 1, 1, 1]).astype('float32')
    mg_y = np.tile(mg_y[None, None, :, :], [batch_size, 1, 1, 1]).astype('float32')
    mesh = torch.cat([torch.tensor(mg_x).to(device), torch.tensor(mg_y).to(device)], 1)
    return mesh


DISTANCE_THRESH_CLEAR = 2
img = imread(PATH + 'train_images/ID_8a6e65317' + '.jpg')
IMG_SHAPE = img.shape

def convert_3d_to_2d(x, y, z, fx=2304.5479, fy=2305.8757, cx=1686.2379, cy=1354.9849):
    # stolen from https://www.kaggle.com/theshockwaverider/eda-visualization-baseline
    return x * fx / z + cx, y * fy / z + cy


def optimize_xy(r, c, x0, y0, z0):
    def distance_fn(xyz):
        x, y, z = xyz
        x, y = convert_3d_to_2d(x, y, z0)
        y, x = x, y
        x = (x - IMG_SHAPE[0] // 2) * IMG_HEIGHT / (IMG_SHAPE[0] // 2) / MODEL_SCALE
        x = np.round(x).astype('int')
        y = (y + IMG_SHAPE[1] // 6) * IMG_WIDTH / (IMG_SHAPE[1] * 4 / 3) / MODEL_SCALE
        y = np.round(y).astype('int')
        return (x - r) ** 2 + (y - c) ** 2

    res = minimize(distance_fn, [x0, y0, z0], method='Powell')
    x_new, y_new, z_new = res.x
    return x_new, y_new, z0


def clear_duplicates(coords):
    for c1 in coords:
        xyz1 = np.array([c1['x'], c1['y'], c1['z']])
        for c2 in coords:
            xyz2 = np.array([c2['x'], c2['y'], c2['z']])
            distance = np.sqrt(((xyz1 - xyz2) ** 2).sum())
            if distance < DISTANCE_THRESH_CLEAR:
                if c1['confidence'] < c2['confidence']:
                    c1['confidence'] = -1
    return [c for c in coords if c['confidence'] > 0]


def extract_coords(prediction, threshold=0):
    logits = prediction[0]
    logits -= threshold # adjust logits by a threshold
    regr_output = prediction[1:]
    points = np.argwhere(logits > 0)
    col_names = sorted(['x', 'y', 'z', 'yaw', 'pitch_sin', 'pitch_cos', 'roll'])
    coords = []
    for r, c in points:
        regr_dict = dict(zip(col_names, regr_output[:, r, c]))
        coords.append(_regr_back(regr_dict))
        coords[-1]['confidence'] = 1 / (1 + np.exp(-logits[r, c]))
        coords[-1]['x'], coords[-1]['y'], coords[-1]['z'] = optimize_xy(r, c, coords[-1]['x'], coords[-1]['y'],
                                                                        coords[-1]['z'])
    coords = clear_duplicates(coords)
    #print('Length:',len(coords))
    return coords


def coords2str(coords, names=['yaw', 'pitch', 'roll', 'x', 'y', 'z', 'confidence']):
    s = []
    for c in coords:
        for n in names:
            s.append(str(c.get(n, 0)))
    return ' '.join(s)


def load_my_state_dict(model, state_dict):
    own_state = model.state_dict()
    print('Total Module to load {}'.format(len(own_state.keys())))
    print('Total Module from weights file {}'.format(len(state_dict.keys())))
    count = 0
    for name, param in state_dict.items():
        if name not in own_state and name.replace('model.module.','') not in own_state:
            continue
        if isinstance(param, torch.nn.Parameter):
            # backwards compatibility for serialized parameters
            param = param.data
        try:
            own_state[name].copy_(param)
        except KeyError as e:
            try:
                own_state[name.replace('model.module.', '')].copy_(param)
            except RuntimeError as e:
                continue
        count += 1

    print('Load Successful {} / {}'.format(count, len(own_state.keys())))


# use a gaussian kernel to spread this 0-1 mask
def gaussian_kernel(x_total, y_total, x_s, y_s, sigma):
    mask = np.zeros((x_total, y_total))
    mask[x_s, y_s] = 1
    x_index = np.arange(0, x_total)
    y_index = np.arange(0, y_total)
    x_index = (x_index - x_s) ** 2
    y_index = (y_index - y_s) ** 2
    x_index = x_index.reshape(-1, 1)
    dist = x_index + y_index
    # gaussian kernel computation
    dist = np.exp(- dist / (2*(sigma ** 2)))
    return mask, dist


'''
functions from kernel: RB's CenterNet Baseline Pytorch without Dropout
'''


def add_number_of_cars(df):
    """df - train or test"""
    df['numcars'] = [int((x.count(' ')+1)/7) for x in df['PredictionString']]
    return df


def remove_out_image_cars(df):
    def isnot_out(x, y):
        # are x,y coordinates within boundaries of the image
        return (x >= 0) & (x <= IMG_SHAPE[1]) & (y >= 0) & (y <= IMG_SHAPE[0])

    df = add_number_of_cars(df)

    new_str_coords = []
    counter_all_ls = []
    for idx, str_coords in enumerate(df['PredictionString']):
        coords = str2coords(str_coords, names=['id', 'yaw', 'pitch', 'roll', 'x', 'y', 'z'])
        xs, ys = get_img_coords(str_coords)
        counter = 0
        coords_new = []

        for (item, x, y) in zip(coords, xs, ys):
            if isnot_out(x, y):
                coords_new.append(item)
                counter += 1

        new_str_coords.append(coords2str(coords_new, names=['id', 'yaw', 'pitch', 'roll', 'x', 'y', 'z']))
        counter_all_ls.append(counter)

    df['new_pred_string'] = new_str_coords
    df['new_numcars'] = counter_all_ls
    print("num of cars outside image bounds:", df['numcars'].sum() - df['new_numcars'].sum(),
          "out of all", df['numcars'].sum(), " cars in train")
    df = df[['ImageId', 'new_pred_string']]
    df.rename(columns={'new_pred_string': 'PredictionString'}, inplace=True)
    return df


def save_submission_file(test, save_dir, predictions, threshold, name=''):
    test['PredictionString'] = predictions
    test.to_csv(save_dir + '/predictions_{}_{}.csv'.format(name, threshold), index=False)
    test = add_number_of_cars(test)
    avg_cars, sum_cars = test.numcars.mean(), test.numcars.sum()
    with open(save_dir + '/stats_{}_{}.txt'.format(name, threshold), 'a+') as f:
        f.write('Average:' + str(avg_cars) + '\n')
        f.write('Total:' + str(avg_cars) + '\n')


def load_checkpoints(args):
    save_dir = args.checkpoint
    models = os.listdir(save_dir)
    models = [i for i in models if 'model' in i]
    models = sorted(models)
    models = models[-1]

    start_epoch = models.replace('.pth','')[models.find('_') + 1:]
    start_epoch = int(start_epoch)
    model = torch.load(os.path.join(save_dir, models))
    if isinstance(model, nn.DataParallel):
        model = model.module
    print('Load checkpoints', models)
    return model, start_epoch