def __init__(self, config):
self.config = config
subdir = datetime.strftime(datetime.now(), '%Y%m%d-%H%M%S')
self.output_dir = os.path.join(config['output_dir'], subdir)
self.model_dir = os.path.join(self.output_dir, 'models')
self.log_dir = os.path.join(self.output_dir, 'log')
self.checkpoint_dir = os.path.join(self.output_dir, 'checkpoints')
self.debug_dir = os.path.join(self.output_dir, 'debug')
check_folders([self.output_dir, self.model_dir, self.log_dir, self.checkpoint_dir, self.debug_dir])
self.val_log = os.path.join(self.output_dir, 'val_log.txt')
self.batch_size = config['batch_size']
self.gpu_num = config['gpu_num']
if self.batch_size % self.gpu_num != 0:
raise ValueError('batch_size must be a multiple of gpu_num')
self.image_size = config['image_size']
self.epoch_num = config['epoch_num']
self.step_per_epoch = config['step_per_epoch']
self.val_freq = config['val_freq']
self.val_data = config['val_data']
self.val_bn_train = config['val_bn_train']
# for k, v in config['val_data'].items():
# self.val_data[k] = load_bin(v, self.image_size)
# imgs = self.val_data[k][0]
# np.save(os.path.join(self.debug_dir, k+'.npy'), imgs[:100])
with open(os.path.join(self.output_dir, 'config.yaml'), 'w') as f:
f.write(yaml.dump(self.config))
def background_extraction(self, cam, frame_num, debug=False, save=False):
check_folders("./mask/{}".format(self.config.strFolderName))
img_color, img_depth = self.get_rgbd(cam, frame_num)
next_mask, num_mask = self.connected_comp_labeling(
cam, img_depth, debug=False)
result = None
prev_mask = self.current_mask[cam].astype(np.uint8)
comparison = np.array(
[compute_img_diff(label, prev_mask, next_mask)
for label in np.unique(next_mask)])
result = np.zeros(next_mask.shape)
result[next_mask == np.argmax(comparison)] = 1 # 255
# for some of the filtering is inverting the expected result
non_zero_size = np.count_nonzero(result)
zero_size = result.shape[0] * result.shape[1] - non_zero_size
if zero_size < non_zero_size:
result = 1 - result
result[result == 1] = 255 # creating the mask for given frame
if save:
self.masks[cam].append(result.astype(np.uint8))
mask = np.dstack((result, result, result)).astype(np.uint8)
clr = np.bitwise_and(img_color, mask)
if debug:
showImageSet([self.current_mask[cam], clr, convert_depth_2_rgb(
img_depth), mask], ["current_mask ", "clr", "depth", "mask"])
self.current_mask[cam] = result.astype(np.uint8)
def save_image_label(self, images, labels, step):
save_dir = os.path.join(self.debug_dir, 'image_by_label')
for i in range(len(labels)):
if(labels[i] < 10):
cur_save_dir = os.path.join(save_dir, str(labels[i]))
check_folders(cur_save_dir)
misc.imsave(os.path.join(cur_save_dir, '%d_%d.jpg' % (step, i)), images[i])
def test_check_folder():
path = "../log/gan/"
check_folders(path, verbose=True, save=True)
assert os.path.exists(path)
assert os.path.exists("../log")
os.remove("../log/gan")
os.remove("../log/")
def save(self, save_path, index, train, label, bk=False):
save_train_folder_path, save_target_folder_path = os.path.join(
save_path, "train"), os.path.join(save_path, "target")
check_folders(save_train_folder_path)
check_folders(save_target_folder_path)
if bk:
self.datasets_bk.append((train, label))
else:
self.datasets.append((train, label))
def __init__(
self,
title,
model_name,
img_shape=(
256,
256,
3),
epochs=100,
batch_size=2,
validation_size=.2,
white_bk=False,
reverse_norm=False):
'''data_model
Loading all of the image from `../data` to self.data
self.data['X']: front image
self.data['y']: back image
'''
self.title = title
self.model_name = model_name
self.batch_size = batch_size
self.reverse_norm = reverse_norm
self.configs = streams_config()
self.nb_epochs = epochs
self.img_shape = img_shape
self.read_list_of_imgs = lambda x: np.array(list(map(vectorized_read_img, x)))
self.weight_path = "../log/" + self.title + "/"
check_folders(self.weight_path)
self.trained_weight_path = os.path.join(
self.weight_path, "%s.h5" %
self.model_name)
self.white_bk = white_bk
# image with max and min with rgb as 0 to 255
self.max = 255.0
self.min = 0.0
# total number of images through out all of the directories
self.total_imgs = self.configs.total_imgs
self.data = {
"X": np.array([]),
"y": np.array([])
}
self.load_datas(validation_size=validation_size)
def __init__(
self,
flag="deconv",
epoch=100000,
img_shape=[256, 256, 3],
learning_rate=1e-3,
white_bk=True,
name="gan_unet_model",
loss=[
'categorical_crossentropy', # generator
l1_loss, # K_dcgan
'binary_crossentropy' #dis criminator
]):
self.loss = loss
assert (flag == "deconv") or (
flag == "upsampling"), "Only support flag for deconv or upsampling"
self.flag = flag
self.img_dim = img_shape
bk = "bk"
if white_bk:
bk = "wh"
data_model.__init__(self,
name + "bk%s_lr_%s_img_dim%s_loss_%s" %
(bk, learning_rate, img_shape[0], loss),
"DCGAN",
img_shape=img_shape,
epochs=epoch)
assert len(loss) == 3
# training params
self.patch_size = [64, 64]
self.n_batch_per_epoch = self.batch_size * 100
self.learning_rate = learning_rate
self.losses = loss
# init all need dir and model
self.build()
self.disc_weights_path = os.path.join(self.weight_path,
"disc_weight_epoch.h5")
self.gen_weights_path = os.path.join(self.weight_path,
"gen_weight_epoch.h5")
self.DCGAN_weights_path = os.path.join(self.weight_path,
"DCGAN_weight_epoch.h5")
check_folders(self.weight_path)
def train(
model_type,
path_args,
training_args,
model_args,
output_args,
load_dataset_args,
checkpoint_args,
extension_args,
):
"""
---------------------------------------------
Input: None
Output: None
Run the test harness for evaluating a model
---------------------------------------------
"""
pretrained = training_args.pop("pretrained")
results_folder = training_args.pop("results_folder")
data_format = load_dataset_args["data_format"]
checkpoint_path = checkpoint_args["checkpoint_path"]
paths = utils.get_paths(**path_args)
utils.check_folders(paths, **extension_args)
keras.backend.set_image_data_format(data_format)
callbacks = []
callbacks.append(utils.get_early_stopping_callback())
callbacks.append(utils.get_tensorboard_directory_callback())
callbacks.append(utils.get_checkpoint_callback(checkpoint_path))
if pretrained:
model = keras.models.load_model(checkpoint_path)
elif model_type == "unet":
model = unet.define_model(output_args, **model_args)
train, val = datagen.load_dataset(paths, load_dataset_args)
history = model.fit_generator(train,
validation_data=val,
callbacks=callbacks,
**training_args)
return (history, model)
def __init__(
self ,
flag="",
epochs=100 ,
img_shape = [ 256 ,256 ,3 ],
learning_rate=1.e-3,
white_bk=True,
batch_size = 20,
name="deep_conv_generative_model",
loss = ["l1",l1_loss]):
bk = "bk"
if white_bk:
bk = "wh"
super.__init__(
self,
"_%s_lr_%s_imgdim%s_loss_%s" %
(bk,
learning_rate,
img_shape[0],
loss),
"DCGAN",
epochs=epochs,
batch_size=batch_size,
white_bk=True)
self.learning_rate = learning_rate
self.generator_loss , self.dc_loss , self.discriminator_loss = \
tuple(loss)
self.patch_size = [64, 64]
self.n_batch_per_epoch = self.batch_size * 100
# init all need dir and model
self.build(self.img_shape)
self.disc_weights_path = os.path.join(
self.weight_path, "disc_weight_epoch.h5")
self.gen_weights_path = os.path.join(
self.weight_path, "gen_weight_epoch.h5")
self.DCGAN_weights_path = os.path.join(
self.weight_path, "DCGAN_weight_epoch.h5")
check_folders(self.weight_path)
def test_img(self):
'''Test the progress of the model during training '''
# pick a random index
idx = rnd.choice([i for i in range(0, len(self.data['X']))])
X, y = self.get_data(idx) # normalized images
self.load()
X_pred = self.model.predict(np.array([X]))
if self.reverse_norm:
X = image.array_to_img(
tanh_inverse_normalization(
X, self.max, self.min))
y = image.array_to_img(
tanh_inverse_normalization(
y, self.max, self.min))
X_pred = image.array_to_img(
tanh_inverse_normalization(
X_pred[0], self.max, self.min))
else:
X = image.array_to_img(
inverse_normalization(
X, self.max, self.min))
y = image.array_to_img(
inverse_normalization(
y, self.max, self.min))
X_pred = image.array_to_img(
inverse_normalization(
X_pred[0], self.max, self.min))
suffix = "End_test"
result = np.hstack((X, y, X_pred))
check_folders("../figures/%s" % (self.title))
plt.imshow(result)
plt.axis("off")
plt.show()
plt.savefig(
"../figures/%s/current_batch_%s.png" %
(self.title, suffix))
def unzip(self, folder):
''' unpack numpy.ndarray to collections of png images train and label'''
data_dir = "./data/%s" % folder
if os.path.isfile(data_dir + "/images.npy"):
check_folders(data_dir + "/train")
check_folders(data_dir + "/target")
data = np.load(data_dir + "/images.npy")
@multi_threads_wrapper(list(chunks(data, 100)))
def save_unzip_imgs(*args):
data, iteration = args
for frame_num, (X, y) in enumerate(data):
train_img = data_dir + \
"/train/train%s.png" % str(iteration + frame_num)
label = data_dir + \
"/target/target%s.png" % str(iteration + frame_num)
cv2.imwrite(train_img, X)
cv2.imwrite(label, y)
save_unzip_imgs()
else:
print("Pre-ziped image was not found ina%s" % data_dir)
import torch
import torch.nn as nn
from torchvision.utils import make_grid
import torch.optim as optim
import torchvision
import torch.nn.functional as F
import torchvision.transforms as transforms
from model import Generator, Discriminator
from utils import showImage, weights_init, check_folders
check_folders()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
tf = transforms.Compose([
transforms.Resize(64),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = torchvision.datasets.CIFAR10(root='./data',
train=True,
download=True,
transform=tf)
#testset = torchvision.datasets.CIFAR10(root = './data', train = False, download = True,
# transform = tf)
