merged two branches to allow training of simple and GNBG model

base_test.py

@@ -0,0 +1,39 @@
+import os
+import sys
+import h5py
+import numpy as np
+from models.model import *
+import tensorflow as tf
+
+source_dir = "results/sparse100_base"
+
+params = dict()
+
+sess = tf.Session()
+
+z1_dim = 128
+z2_dim = 3
+x_dim = 3
+d_dim = 128
+num_objs = 1000
+
+data_file = 'data/old_data/data_100.h5'
+with h5py.File(data_file, 'r') as f:
+    images = np.array(f['test/images'], dtype=np.float32)
+    sparse = np.array(f['test/sparse'], dtype=np.float32)
+    depths = np.array(f['test/depths'], dtype=np.float32)
+
+input_imgs = tf.concat([images, tf.expand_dims(sparse, axis=3)], axis=3)
+net = Network(input_imgs, params)
+est_maps = tf.reshape(net.forward(), [-1, 200, 200])
+
+loss = tf.reduce_mean(tf.abs((est_maps - depths)))
+
+sess.run(tf.global_variables_initializer())   
+
+
+saver = tf.train.Saver()
+saver.restore(sess, source_dir + '/model.ckpt')
+
+test_error = loss.eval(session=sess)
+print(test_error)

configs/local_global.yml

@@ -0,0 +1,7 @@
+type: "local_global"
+data_file: "data/old_data/data_01.h5"
+batch_size: 32
+num_iters: 10000
+optimizer: "sgd"
+learning_rate: 0.001
+save_output: False

configs/simple.yml

@@ -1,6 +1,7 @@
-data_file: "data/data_01.h5"
+type: "simple"
+data_file: "data/old_data/data_01.h5"
 batch_size: 32
-
 num_iters: 10000
-latent_dim: 128
+optimizer: "sgd"
 learning_rate: 0.001
+save_output: False

main.py

@@ -2,6 +2,8 @@
 import yaml
 
 from trainers.base_trainer import BaseTrainer
+from trainers.simple_trainer import SimpleTrainer
+from trainers.local_global_trainer import LocalGlobalTrainer
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--config", type=str, help='path to config file')
@@ -14,5 +16,8 @@
 with open("configs/" + args.config, 'r') as f:
     params = yaml.load(f, Loader=yaml.FullLoader)
 
-trainer = BaseTrainer(params)
+if params["type"] == "simple":
+    trainer = SimpleTrainer(params)
+elif params["type"] == "local_global":
+    trainer = LocalGlobalTrainer(params)
 trainer.train()

models/baseline_avg.py

@@ -0,0 +1,47 @@
+import os
+import rasterio
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+import scipy.misc
+from PIL import Image
+from sklearn.model_selection import train_test_split
+from sklearn.cluster import KMeans
+from sklearn.feature_extraction.image import img_to_graph
+from scipy.interpolate import interp2d
+from itertools import combinations_with_replacement
+from collections import defaultdict
+import random
+import h5py
+
+data_file = "../data/data_01.h5"
+debug = False
+
+data = h5py.File(data_file, 'r')
+
+test_images = data["test/images"]
+test_sparse = data["test/sparse"]
+test_depths = data["test/depths"]
+
+
+total = 0.0           
+for i in tqdm(range(test_images.shape[0])):
+    image = 1/3 * test_images[i, :, :, 0] + 1/3 * test_images[i, :, :, 1] + 1/3 * test_images[i, :, :, 2]
+    sparse = test_sparse[i, :, :]
+    depth = np.zeros(sparse.shape)
+    gt = test_depths[i, :, :]
+
+    if np.count_nonzero(sparse) != 0:
+        avg_depth = sparse.sum() / np.count_nonzero(sparse)
+    else:
+        avg_depth = sparse.sum()
+    depth[:, :] = avg_depth
+    if debug:
+        plt.imsave('depth' + str(i) + '.jpg', depth[:, :])
+        plt.imsave('gt' + str(i) + '.jpg', gt[:, :])
+        plt.imsave('img' + str(i) + '.jpg', image)
+        plt.imsave('sparse' + str(i) + '.jpg', sparse)
+    error = (np.abs((depth-gt))).mean()
+    #print(error)
+    total = total + error
+print("Total error:", total/test_images.shape[0])

models/baseline_identity.py

@@ -0,0 +1,34 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+import scipy
+import h5py
+
+data_file = "../data/data_01.h5"
+debug = False
+
+data = h5py.File(data_file, 'r')
+
+test_images = data["test/images"]
+test_sparse = data["test/sparse"]
+test_depths = data["test/depths"]
+
+
+total = 0.0           
+for i in tqdm(range(test_images.shape[0])):
+    image = 1/3 * test_images[i, :, :, 0] + 1/3 * test_images[i, :, :, 1] + 1/3 * test_images[i, :, :, 2]
+    sparse = test_sparse[i, :, :]
+    depth = np.zeros(sparse.shape)
+    gt = test_depths[i, :, :]
+
+    depth = sparse
+
+    if debug:
+        plt.imsave('depth' + str(i) + '.jpg', depth[:, :])
+        plt.imsave('gt' + str(i) + '.jpg', gt[:, :])
+        plt.imsave('img' + str(i) + '.jpg', image)
+    error = (np.abs((depth-gt))).mean()
+    #print(error)
+    total = total + error
+print("Total error:", total/test_images.shape[0])

models/baseline_segmentation.py

@@ -0,0 +1,58 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+from sklearn.model_selection import train_test_split
+from sklearn.cluster import KMeans
+from sklearn.feature_extraction.image import img_to_graph
+from scipy.interpolate import interp2d
+import h5py
+
+data_file = "../data/data_01.h5"
+debug = False
+
+data = h5py.File(data_file, 'r')
+
+test_images = data["test/images"]
+test_sparse = data["test/sparse"]
+test_depths = data["test/depths"]
+
+
+total = 0.0           
+for i in tqdm(range(test_images.shape[0])):
+    image = 1/3 * test_images[i, :, :, 0] + 1/3 * test_images[i, :, :, 1] + 1/3 * test_images[i, :, :, 2]
+    sparse = test_sparse[i, :, :]
+    depth = np.zeros(sparse.shape)
+    gt = test_depths[i, :, :]
+
+    #x = []
+    #y = []
+    #z = []
+    #for row in range(200):
+    #    for col in range(200):
+    #        entry = sparse[row, col]
+    #        if entry != 0:
+    #            x.append(row)
+    #            y.append(col)
+    #            z.append(entry)
+    #f = interp2d(x, y, z)
+    #x = range(0, 200)
+    #y = range(0, 200)
+    #depth = f(x, y)
+    #print(depth.shape)
+    x = np.reshape(image, [-1, 1])
+    print("Fitting k-means...")
+    labels = KMeans(n_clusters=5).fit_predict(x)
+    label_img = np.reshape(labels, [200, 200])
+    for label in labels:
+        mask = label_img == label
+        avg_depth = sparse[mask].sum() / np.count_nonzero(sparse[mask])
+        depth[mask] = avg_depth
+    if debug:
+        plt.imsave('depth' + str(i) + '.jpg', depth[:, :])
+        plt.imsave('gt' + str(i) + '.jpg', gt[:, :])
+        plt.imsave('labels' + str(i) + '.jpg', label_img)
+        plt.imsave('limg' + str(i) + '.jpg', image)
+    error = (np.abs((depth-gt))).mean()
+    print(error)
+    total = total + error
+print("Total error:", total/test_images.shape[0])