step2_segm_vote_gmm.py

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import cv2
import os
import gco
import argparse
import numpy as np
import cPickle as pkl

from glob import glob
from scipy import signal

from util.labels import LABELS_REDUCED, LABEL_COMP, LABELS_MIXTURES, read_segmentation
from sklearn.mixture import GaussianMixture


def edges_seams(seams, tex_res, edge_idx):
    edges = np.zeros((0, 2), dtype=np.int32)

    for _, e0, _, e1 in seams:
        idx0 = np.array(edge_idx[e0][0]) * tex_res + np.array(edge_idx[e0][1])
        idx1 = np.array(edge_idx[e1][0]) * tex_res + np.array(edge_idx[e1][1])

        if len(idx0) and len(idx1):
            if idx0.shape[0] < idx1.shape[0]:
                idx0 = cv2.resize(idx0.reshape(-1, 1), (1, idx1.shape[0]), interpolation=cv2.INTER_NEAREST)
            elif idx0.shape[0] > idx1.shape[0]:
                idx1 = cv2.resize(idx1.reshape(-1, 1), (1, idx0.shape[0]), interpolation=cv2.INTER_NEAREST)

            edges_new = np.hstack((idx0.reshape(-1, 1), idx1.reshape(-1, 1)))
            edges = np.vstack((edges, edges_new))

    edges = np.sort(edges, axis=1)

    return edges[:, 0], edges[:, 1]


def main(unwrap_dir, segm_out_file, gmm_out_file):
    iso_files = np.array(sorted(glob(os.path.join(unwrap_dir, '*_unwrap.jpg'))))
    segm_files = np.array(sorted(glob(os.path.join(unwrap_dir, '*_segm.png'))))
    vis_files = np.array(sorted(glob(os.path.join(unwrap_dir, '*_visibility.jpg'))))

    iso_mask = cv2.imread('assets/tex_mask_1000.png', flags=cv2.IMREAD_GRAYSCALE) / 255.
    iso_mask = cv2.resize(iso_mask, (1000, 1000), interpolation=cv2.INTER_NEAREST)

    voting = np.zeros((1000, 1000, len(LABELS_REDUCED)))

    gmms = {}
    gmm_pixels = {}

    for color_id in LABELS_REDUCED:
        gmms[color_id] = GaussianMixture(LABELS_MIXTURES[color_id])
        gmm_pixels[color_id] = []

    for frame_file, segm_file, vis_file in zip(iso_files, segm_files, vis_files):
        print('extract from {}...'.format(os.path.basename(frame_file)))

        frame = cv2.cvtColor(cv2.imread(frame_file), cv2.COLOR_BGR2HSV) / 255.
        tex_segm = read_segmentation(segm_file)
        tex_weights = 1 - cv2.imread(vis_file) / 255.
        tex_weights = np.sqrt(tex_weights)

        for i, color_id in enumerate(LABELS_REDUCED):
            if color_id != 'Unseen' and color_id != 'BG':
                where = np.all(tex_segm == LABELS_REDUCED[color_id], axis=2)
                voting[where, i] += tex_weights[where, 0]
                gmm_pixels[color_id].extend(frame[where].tolist())

    for color_id in LABELS_REDUCED:
        if gmm_pixels[color_id]:
            print('GMM fit {}...'.format(color_id))
            gmms[color_id].fit(np.array(gmm_pixels[color_id]))

    for i, color_id in enumerate(LABELS_REDUCED):
        if color_id == 'Unseen' or color_id == 'BG':
            voting[:, i] = -10

    voting[iso_mask == 0] = 0
    voting[iso_mask == 0, 0] = 1

    unaries = np.ascontiguousarray((1 - voting / len(iso_files)) * 10)
    pairwise = np.ascontiguousarray(LABEL_COMP)

    seams = np.load('assets/basicModel_seams.npy')
    edge_idx = pkl.load(open('assets/basicModel_edge_idx_1000.pkl', 'rb'))

    dr_v = signal.convolve2d(iso_mask, [[-1, 1]])[:, 1:]
    dr_h = signal.convolve2d(iso_mask, [[-1], [1]])[1:, :]

    where_v = iso_mask - dr_v
    where_h = iso_mask - dr_h

    idxs = np.arange(1000 ** 2).reshape(1000, 1000)
    v_edges_from = idxs[:-1, :][where_v[:-1, :] == 1].flatten()
    v_edges_to = idxs[1:, :][where_v[:-1, :] == 1].flatten()
    h_edges_from = idxs[:, :-1][where_h[:, :-1] == 1].flatten()
    h_edges_to = idxs[:, 1:][where_h[:, :-1] == 1].flatten()

    s_edges_from, s_edges_to = edges_seams(seams, 1000, edge_idx)

    edges_from = np.r_[v_edges_from, h_edges_from, s_edges_from]
    edges_to = np.r_[v_edges_to, h_edges_to, s_edges_to]
    edges_w = np.r_[np.ones_like(v_edges_from), np.ones_like(h_edges_from), np.ones_like(s_edges_from)]

    gc = gco.GCO()
    gc.create_general_graph(1000 ** 2, pairwise.shape[0], True)
    gc.set_data_cost(unaries.reshape(1000 ** 2, pairwise.shape[0]))

    gc.set_all_neighbors(edges_from, edges_to, edges_w)
    gc.set_smooth_cost(pairwise)
    gc.swap(-1)

    labels = gc.get_labels().reshape(1000, 1000)
    gc.destroy_graph()

    segm_colors = np.zeros((1000, 1000, 3), dtype=np.uint8)

    for i, color_id in enumerate(LABELS_REDUCED):
        segm_colors[labels == i] = LABELS_REDUCED[color_id]

    cv2.imwrite('{}'.format(segm_out_file), segm_colors[:, :, ::-1])
    pkl.dump(gmms, open(gmm_out_file, 'wb'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        'unwrap_dir',
        type=str,
        help="Directory that contains unwrap files")

    parser.add_argument(
        'segm_out_file',
        type=str,
        help="Output file for segmentation")

    parser.add_argument(
        'gmm_out_file',
        type=str,
        help="Output file for GMMs")

    args = parser.parse_args()

    main(args.unwrap_dir, args.segm_out_file, args.gmm_out_file)