# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

from pylab import *
import os
import sys
from sklearn.mixture import GaussianMixture
import pickle
from PIL import Image
import transforms3d

dataset_dir = 'rgbd_dataset_freiburg3_long_office_household'
gmm_dir = 'gmm_qqvga_mine_patch'
full_gmm_dir = os.path.join(dataset_dir,gmm_dir)

files = []
names = []
for fl in sorted(os.listdir(full_gmm_dir)):
    name,ext = os.path.splitext(fl)

    with open(os.path.join(full_gmm_dir,fl),'rb') as fp:
        clf = pickle.load(fp)
        files.append(clf)
    names.append(name)
initial_state_str = '-0.6832 2.6909 1.7373 0.0003 0.8617 -0.5072 -0.0145'
initial_state = np.array([float(_) for _ in initial_state_str.split(' ')])


import scipy.optimize as opt
#len(files)-1
results_trans = []
ttimes = 0.0
for i in range(len(files)-1):
    from_fl = files[i]
    to_fl = files[i+1] #i+1

    from_params = from_fl._get_parameters()
    to_params = to_fl._get_parameters()

    def loss_gmm_to_gmm(x):
        qs = x[:4]
        ts = x[4:]
        qs = qs/np.linalg.norm(qs)
        Ms = transforms3d.quaternions.quat2mat(qs)
        
        y = from_params[1] - (Ms @ to_params[1].T).T - ts
        sigma = from_params[2] + Ms @ to_params[2] @ Ms.T
        inv_sigma = np.array([np.linalg.inv(s) for s in sigma])
        dist = y[:,None,:] @ inv_sigma @ y[:,:,None]
        f = np.sqrt(np.linalg.det(inv_sigma))* from_params[0] * to_params[0] * np.exp(-0.5 * dist)
        return -f.sum()
    def loss_gmm_to_gmm2(x):
        qs = x[:4]
        ts = x[4:]
        qs = qs/np.linalg.norm(qs)
        Ms = transforms3d.quaternions.quat2mat(qs)
        
        a = from_params[1] - ts
        b = (Ms @ to_params[1].T).T
        #y2 = np.repeat(a,b.shape[0],0)-np.tile(b,(a.shape[0],1))
        #y2 = y2.sum(1).reshape((a.shape[0],a.shape[0]))
        y = a - b[:, None]

        sigma = from_params[2] + Ms @ to_params[2] @ Ms.T
        inv_sigma = np.array([np.linalg.inv(s) for s in sigma])
        dist = y.T @ inv_sigma @ y
        print(dist.shape)
        f = np.sqrt(np.linalg.det(inv_sigma))* from_params[0] * to_params[0] * np.exp(-0.5 * dist)
        return -f.sum()
    def loss_gmm_to_gmm_raw(x):
        qs = x[:4]
        ts = x[4:]
        qs = qs/np.linalg.norm(qs)
        Ms = transforms3d.quaternions.quat2mat(qs)
        total = 0
        for i in range(100):
            for j in range(100):
                y = from_params[1][i:i+1].T - (Ms @ to_params[1][j:j+1].T) - ts[:,None]
                sigma = from_params[2][i] + Ms @ to_params[2][j] @ Ms.T
                inv_sigma = np.linalg.pinv(sigma)
                dist = y.T @ inv_sigma @ y
                dist = dist.sum()
                f = np.sqrt(np.linalg.det(inv_sigma))* from_params[0][i] * to_params[0][j] * np.exp(-0.5 * dist)
                total += f

        return -total
    def loss_gmm_to_gmm_raw2(x):
        qs = x[:4]
        ts = x[4:]
        qs = qs/np.linalg.norm(qs)
        Ms = transforms3d.quaternions.quat2mat(qs)
        total = 0
        for i in range(100):
            y = from_params[1][i:i+1].T - (Ms @ to_params[1].T) - ts[:,None]
            y = y.T
            sigma = from_params[2][i].T + Ms @ to_params[2] @ Ms.T
            inv_sigma = np.linalg.inv(sigma)
            dist = np.squeeze(y[:,None,:] @ inv_sigma @ y[:,:,None])
            f = np.sqrt(np.linalg.det(inv_sigma))* from_params[0][i] * to_params[0] * np.exp(-0.5 * dist)
            total += f.sum()

        return -total
    def loss_gmm_to_gmm_raw3(x):
        qs = x[:4]
        ts = x[4:]
        qs = qs/np.linalg.norm(qs)
        Ms = transforms3d.quaternions.quat2mat(qs)
        total = 0
        a = from_params[1] - ts
        b = (Ms @ to_params[1].T).T
        y = a - b[:, None]
        y = y.reshape((-1,3))

        sigma = from_params[2] + (Ms @ to_params[2] @ Ms.T)[:,None]
        sigma = sigma.reshape((-1,3,3))
        inv_sigma = np.linalg.inv(sigma)

        dist = np.squeeze(y[:,None,:] @ inv_sigma @ y[:,:,None])
        dist = dist.reshape((a.shape[0],b.shape[0]))
        scaler = 1#np.sqrt(np.linalg.det(inv_sigma)).reshape((a.shape[0],b.shape[0]))
        f = scaler * from_params[0][None,:] * to_params[0][:,None] * np.exp(-0.5 * dist)
        total += f.sum()
        return -total
    #print(loss_gmm_to_gmm(np.array([1,0,0,0,0,0,0])))
    #print(loss_gmm_to_gmm_raw(np.array([1,0,0,0,0,0,0])))
    #print(loss_gmm_to_gmm_raw2(np.array([1,0,0,0,0,0,0])))

    #print(loss_gmm_to_gmm2(np.array([1,0,0,0,0,0,0])))
    cg_cnt = 0
    t1 = time.time()
    res = opt.minimize(loss_gmm_to_gmm_raw3,np.array([1,0,0,0,0,0,0]),method=None)
    #res2 = opt.minimize(loss_gmm_to_gmm_raw3,np.array([1,0,0,0,0,0,0]),method='CG')
    t2 = time.time()
    print(i,t2-t1,ttimes/(i+1))
    #raise
    ttimes += (t2 - t1)
    x = res.x
    qe = x[:4]
    qe = qe/np.linalg.norm(qe)
    te = x[4:]
    qe,te
    results_trans.append((te,qe))
    res.fun
print('total_time',ttimes)

# +

print(loss_gmm_to_gmm_raw(np.array([1,0,0,0,0,0,0])))
print(loss_gmm_to_gmm_raw3(np.array([1,0,0,0,0,0,0])))

# -

results_trans

import pickle
with open('patch2_nodet.pkl','wb') as fp:
    pickle.dump(results_trans,fp)
transforms3d.quaternions.quat2mat(qe)

name_times = [float(_) for _ in names]


state = [0,0,0,1,0,0,0]
#state = [-0.6832,2.6909,1.7373,0.0003,0.8617,-0.5072,-0.0145]
states = []
states.append(state)
log_writer = []
p_state = np.copy(state)
w = p_state[3]
p_state[3] = p_state[-1]
p_state[-1] = w
log_writer.append(names[0] + ' ' + ' '.join([str(_) for _ in p_state]))
for i,d in enumerate(results_trans):
    te,qe = d
    Rt = transforms3d.quaternions.quat2mat(qe)
    state_t = state[:3]
    state_q = state[3:]
    RW = transforms3d.quaternions.quat2mat(state_q)
    RW =    Rt @ RW
    new_state_t = state_t -  te @ RW.T
    new_state_q = transforms3d.quaternions.mat2quat(RW)

    print(new_state_t,new_state_q)

    state = np.hstack([new_state_t,new_state_q])
    w = new_state_q[0]
    new_state_q[0] = new_state_q[3]
    new_state_q[3] = w
    log_writer.append(' '.join([names[i+1],' '.join([str(_) for _ in new_state_t]),' '.join([str(_) for _ in new_state_q])]))

len(log_writer)

with open('gmm_res_patch2_nodet.txt','wt') as fp:
    for line in log_writer:
        fp.write(line +'\n')

# # visualize an example

# %matplotlib notebook
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource

means,covars,weights = files[0].means_,files[0].covariances_,files[0].weights_
fig = plt.figure()
ax = Axes3D(fig)
colors = plt.cm.Pastel1(np.arange(means.shape[0])/(means.shape[0])) #plt.rcParams['axes.prop_cycle'].by_key()['color']
np.random.shuffle(colors)
for i in range(means.shape[0]):
    orig_i = i
    mean,covar = means[i],covars[i]
    u,s,vt = np.linalg.svd(covar)
    coefs = (.002, .002, .002)  # Coefficients in a0/c x**2 + a1/c y**2 + a2/c z**2 = 1 
    # Radii corresponding to the coefficients:
    rx, ry, rz = 1.7*np.sqrt(s)#s#1/np.sqrt(coefs)
    
    R_reg = vt.T @ np.diag([1,1,np.linalg.det(vt.T @ u.T)]) @ u.T
    
    #print(eigs)
    # Set of all spherical angles:
    u = np.linspace(0, 2 * np.pi, 10)
    v = np.linspace(0, np.pi, 10)

    # Cartesian coordinates that correspond to the spherical angles:
    # (this is the equation of an ellipsoid):
    x = rx * np.outer(np.cos(u), np.sin(v)) #+ mean[0]
    y = ry * np.outer(np.sin(u), np.sin(v)) #+ mean[1]
    z = rz * np.outer(np.ones_like(u), np.cos(v)) #+ mean[2]
    
    for i in range(len(x)):
        for j in range(len(x)):
            x[i,j],y[i,j],z[i,j] = np.dot([x[i,j],y[i,j],z[i,j]], vt) + mean    
    # Plot:
    res = ax.plot_surface(x,y,z,  color=colors[orig_i],shade=True,linewidth=0.0)

# +
means,covars,weights = files[10].means_,files[10].covariances_,files[10].weights_
fig = plt.figure()
ax = Axes3D(fig)
for i in range(means.shape[0]):
    orig_i = i
    mean,covar = means[i],covars[i]
    u,s,vt = np.linalg.svd(covar)
    # Radii corresponding to the coefficients:
    rx, ry, rz = 3*np.sqrt(s)#s#1/np.sqrt(coefs)
    
    R_reg = vt.T @ np.diag([1,1,np.linalg.det(vt.T @ u.T)]) @ u.T
    
    #print(eigs)
    # Set of all spherical angles:
    u = np.linspace(0, 2 * np.pi, 10)
    v = np.linspace(0, np.pi, 10)

    # Cartesian coordinates that correspond to the spherical angles:
    # (this is the equation of an ellipsoid):
    x = rx * np.outer(np.cos(u), np.sin(v)) #+ mean[0]
    y = ry * np.outer(np.sin(u), np.sin(v)) #+ mean[1]
    z = rz * np.outer(np.ones_like(u), np.cos(v)) #+ mean[2]
    
    for i in range(len(x)):
        for j in range(len(x)):
            x[i,j],y[i,j],z[i,j] = np.dot([x[i,j],y[i,j],z[i,j]], vt) + mean    
    # Plot:
    res = ax.plot_surface(x,y,z,color='r',shade=True,linewidth=0.0,alpha=np.minimum(1.0,weights[orig_i]*10))
    
means,covars,weights = files[185].means_,files[185].covariances_,files[185].weights_
for i in range(means.shape[0]):
    orig_i = i
    mean,covar = means[i],covars[i]
    u,s,vt = np.linalg.svd(covar)
    # Radii corresponding to the coefficients:
    rx, ry, rz = 3*np.sqrt(s)#s#1/np.sqrt(coefs)
    
    R_reg = vt.T @ np.diag([1,1,np.linalg.det(vt.T @ u.T)]) @ u.T
    
    #print(eigs)
    # Set of all spherical angles:
    u = np.linspace(0, 2 * np.pi, 10)
    v = np.linspace(0, np.pi, 10)

    # Cartesian coordinates that correspond to the spherical angles:
    # (this is the equation of an ellipsoid):
    x = rx * np.outer(np.cos(u), np.sin(v)) #+ mean[0]
    y = ry * np.outer(np.sin(u), np.sin(v)) #+ mean[1]
    z = rz * np.outer(np.ones_like(u), np.cos(v)) #+ mean[2]
    
    for i in range(len(x)):
        for j in range(len(x)):
            x[i,j],y[i,j],z[i,j] = np.dot([x[i,j],y[i,j],z[i,j]], vt) + mean    
    # Plot:
    res = ax.plot_surface(x,y,z,color='b',shade=True,linewidth=0.0,alpha=np.minimum(1.0,weights[orig_i]*10))
ax.view_init(elev=12., azim=-4)
ax.set_xlim3d(-3,3)
ax.set_ylim3d(-3,3)
ax.set_zlim3d(.5,6.5)

# +
import pandas as pd
gt = pd.read_csv('rgbd_dataset_freiburg3_long_office_household-groundtruth.txt',' ',comment='#',header=None,names='timestamp tx ty tz qx qy qz qw'.split(' '))
def row_to_vec(row):
    a = np.array(row)
    return a[1:4], a[4:]
from scipy.spatial.distance import cdist, pdist
time_dists = cdist(np.array(gt.timestamp)[:,None],np.array(name_times)[:,None])
matches = time_dists.argmin(axis=0)

frm = row_to_vec(gt.iloc[matches[10]])
to = row_to_vec(gt.iloc[matches[100]])
matches[10],matches[100]
# -

time_dists[26,10],time_dists[630,185]

frm[1]/np.linalg.norm(frm[1]),to[1]/np.linalg.norm(to[1])

# +
ttrue = frm[0] - to[0]
rf =  transforms3d.quaternions.quat2mat(frm[1]/np.linalg.norm(frm[1]))
rt = transforms3d.quaternions.quat2mat(to[1]/np.linalg.norm(to[1]))

#Rtrue = transforms3d.quaternions.quat2mat(quat)
Rtrue = rt.T @ rf

# +
means,covars,weights = files[10].means_,files[10].covariances_,files[10].weights_
fig = plt.figure()
ax = Axes3D(fig)
for i in range(means.shape[0]):
    orig_i = i
    mean,covar = means[i],covars[i]
    u,s,vt = np.linalg.svd(covar)
    # Radii corresponding to the coefficients:
    rx, ry, rz = 3*np.sqrt(s)#s#1/np.sqrt(coefs)
    
    R_reg = vt.T @ np.diag([1,1,np.linalg.det(vt.T @ u.T)]) @ u.T
    
    #print(eigs)
    # Set of all spherical angles:
    u = np.linspace(0, 2 * np.pi, 12)
    v = np.linspace(0, np.pi, 12)

    # Cartesian coordinates that correspond to the spherical angles:
    # (this is the equation of an ellipsoid):
    x = rx * np.outer(np.cos(u), np.sin(v)) #+ mean[0]
    y = ry * np.outer(np.sin(u), np.sin(v)) #+ mean[1]
    z = rz * np.outer(np.ones_like(u), np.cos(v)) #+ mean[2]
    
    for i in range(len(x)):
        for j in range(len(x)):
            x[i,j],y[i,j],z[i,j] = np.dot([x[i,j],y[i,j],z[i,j]], vt) + mean    
    # Plot:
    res = ax.plot_surface(x,y,z,color='r',shade=True,linewidth=0.0,alpha=np.minimum(1.0,weights[orig_i]*10))
    
means,covars,weights = files[100].means_,files[100].covariances_,files[100].weights_
for i in range(means.shape[0]):
    orig_i = i
    mean,covar = means[i]+ttrue, covars[i] @ Rtrue.T
    u,s,vt = np.linalg.svd(covar)
    # Radii corresponding to the coefficients:
    rx, ry, rz = 3*np.sqrt(s)#s#1/np.sqrt(coefs)
    
    R_reg = vt.T @ np.diag([1,1,np.linalg.det(vt.T @ u.T)]) @ u.T
    
    #print(eigs)
    # Set of all spherical angles:
    u = np.linspace(0, 2 * np.pi, 12)
    v = np.linspace(0, np.pi, 12)

    # Cartesian coordinates that correspond to the spherical angles:
    # (this is the equation of an ellipsoid):
    x = rx * np.outer(np.cos(u), np.sin(v)) #+ mean[0]
    y = ry * np.outer(np.sin(u), np.sin(v)) #+ mean[1]
    z = rz * np.outer(np.ones_like(u), np.cos(v)) #+ mean[2]
    
    for i in range(len(x)):
        for j in range(len(x)):
            x[i,j],y[i,j],z[i,j] = np.dot([x[i,j],y[i,j],z[i,j]], vt) + mean    
    # Plot:
    res = ax.plot_surface(x,y,z,color='b',shade=True,linewidth=0.0,alpha=np.minimum(1.0,weights[orig_i]*10))
ax.view_init(elev=12., azim=-4)
ax.set_xlim3d(-2,2)
ax.set_ylim3d(-2,2)
ax.set_zlim3d(.5,4.5)
# -







StarcoderdataJupyterScriptsDedupFiltered