def create_model(model='', image_shape=[1024, 1024], class_num=9): train_image = fluid.layers.data(name='img', shape=[3] + image_shape, dtype='float32') train_label = fluid.layers.data(name='label', shape=image_shape + [9], dtype='float32') if model == 'unet': predict = unet().model(train_image) if model == 'deeplab_v3p': predict = deeplab_v3p().model(train_image) if model == 'pannet': predict = pannet().model(train_image) if model == 'dense_unet': predict = dense_unet().model(train_image) if model == 'multires_unet': predict = multires_unet().model(train_image) if model == 'deeplabv3p_ours': predict = deeplabv3p_ours().model(train_image) loss = categorical_crossentropy(train_label, predict) iou, out_wrong, out_right = create_iou(predict, train_label, class_num) return predict, loss, iou
def _available_2d_nets(network, window_shape, n_class, preset_model): """Returns the chosen 2D network model adapted for window shape and number of classes.""" available_nets = { 'tiramisu': tiramisu(input_size=(*window_shape, n_class), preset_model=preset_model), 'unet': unet(input_size=(*window_shape, n_class)), } return available_nets.get(network, None)
def get_model(model_type, multi_modal, perform_early_fusion, pe_block, inputA, inputB, cascade=False): if model_type == "unet": return unet(multi_modal, perform_early_fusion, pe_block, inputA, inputB, cascade) elif model_type == "unet++": return unetpp(multi_modal, perform_early_fusion, pe_block, inputA, inputB, cascade)
def get_model(algorithm: str, input_size: int, num_classes: int, loss, channels: int = 3, ): if algorithm == 'fcn_densenet': model, model_name = fcndensenet.fcndensenet(input_size=input_size, num_classes=num_classes, loss=loss, channels=channels) elif algorithm == 'tiramisu': model, model_name = tiramisu.tiramisu(input_size=input_size, num_classes=num_classes, loss=loss, channels=channels) elif algorithm == 'unet': model, model_name = unet.unet(input_size=input_size, num_classes=num_classes, loss=loss, channels=channels) elif algorithm == 'pspnet': model, model_name = pspnet.pspnet(input_size=input_size, num_classes=num_classes, loss=loss, channels=channels) elif algorithm == 'deeplabv3plus': model, model_name = (deeplabv3plus.Deeplabv3(weights=None, input_shape=(input_size, input_size, channels), classes=num_classes, backbone='xception', loss=loss), "deeplabv3plus") else: raise Exception('{} is an invalid algorithm'.format(algorithm)) return model, model_name
def create_model(model='',image_shape=[1024,1024],class_num=9): train_image = fluid.layers.data(name='img', shape=[3] + image_shape, dtype='float32') if model == 'unet': predict = unet().model(train_image) if model == 'deeplab_v3p': predict = deeplab_v3p().model(train_image) if model == 'pannet': predict = pannet().model(train_image) if model == 'dense_unet': predict = dense_unet().model(train_image) if model == 'multires_unet': predict = multires_unet().model(train_image) if model == 'deeplabv3p_ours': predict = deeplabv3p_ours().model(train_image) return predict
def test(FLAGS): ### Parameters of the testing cropShape = (192, 192, 3) test_path = FLAGS.test_path noise_level = FLAGS.noise architecture = FLAGS.architecture filepath = FLAGS.load_weights_path sliding_window = FLAGS.sliding_window # Load model if architecture == "unet": model = unet(cropShape) elif architecture == "wavelet": model = unetWavelet(cropShape) else: raise RuntimeError( 'Unkwown architecture, please choose from "unet" or "wavelet".') model.load_weights(filepath) # Load the images in the test folder test_files = [img_path for img_path in glob(test_path + '/*.png')] for img_path in test_files: image = load_img(img_path) image = img_to_array(image) if sliding_window != 0: image = image[:sliding_window * ((image.shape[0] - cropShape[0]) // sliding_window) + cropShape[0], :sliding_window * ((image.shape[1] - cropShape[1]) // sliding_window) + cropShape[1]] else: image = image[:(image.shape[0] // cropShape[0]) * cropShape[0], :(image.shape[1] // cropShape[1]) * cropShape[1]] # Add noise noise = np.random.normal(0, noise_level, image.shape) noisy_image = image + noise noisy_image /= 255.0 residual_image = np.empty(image.shape) # Divide the image into subimages of size cropShape to match the model's input size. if sliding_window != 0: sum_img = np.zeros(image.shape) count_img = np.zeros(image.shape) for y in range(0, sliding_window + image.shape[0] - cropShape[0], sliding_window): for x in range(0, sliding_window + image.shape[1] - cropShape[1], sliding_window): extract = noisy_image[y:y + cropShape[0], x:x + cropShape[1]] sum_img[y:y + cropShape[0], x:x + cropShape[1]] += model.predict(extract[np.newaxis, ...])[0] count_img[y:y + cropShape[0], x:x + cropShape[1]] += 1 residual_image = sum_img / count_img else: for block_y in range(0, image.shape[0] // cropShape[0]): for block_x in range(0, image.shape[1] // cropShape[1]): extract = noisy_image[block_y * cropShape[0]:(block_y + 1) * cropShape[0], block_x * cropShape[1]:(block_x + 1) * cropShape[1]] residual_image[block_y * cropShape[0]:(block_y + 1) * cropShape[0], block_x * cropShape[1]:(block_x + 1) * cropShape[1]] = model.predict( extract[np.newaxis, ...])[0] # Display the noised and denoised images display_img(noisy_image, residual_image)
def main(argv): data_path = None input_image = None output = None weight_path = None mode=5 drop_rate=0 lab=s.lab classification=False temp=.4 try: opts, args = getopt.getopt(argv,"w:p:b:m:ld:ct:i:o:",["weight-path=", "datapath=",'model=','lab','drop-rate=','input=','output=']) except getopt.GetoptError as error: print(error) print( 'demo.py -w <path to weights file> -p <path to folder of images> OR -i <path to single image> -l <no argument. use if lab should be used>\ -d <amount of dropout used in model> -c <no argument. Use if model is classifier> -t <temperature for annealed mean> -o <output path for images>') sys.exit(2) print("opts", opts) for opt, arg in opts: if opt in ("-w", "--weight-path"): weight_path = arg elif opt in ("--datapath", "-p"): data_path = arg elif opt=='-m': if arg in ('custom','0'): mode = 0 elif arg in ('u','1','unet'): mode = 1 elif arg in ('ende','2'): mode = 2 elif arg in ('richzhang','classende','3'): mode = 3 elif arg in ('colorunet','cu','4'): mode = 4 elif arg in ('mu','5','middle'): mode = 5 elif opt in ('-l','--lab'): lab=True elif opt in ("-d", "--drop-rate"): drop_rate = float(arg) elif opt =='-c': classification=True lab=True elif opt=='-t': temp=float(arg) elif opt in ('-i','--input'): input_image = arg elif opt in ('-o','--output'): output = arg if data_path is None and input_image is None: print('Please select an image or folder') sys.exit() trafo=transforms.Compose([transforms.Grayscale(3 if lab else 1), transforms.Resize((96,96))]) device=torch.device('cuda' if torch.cuda.is_available() else 'cpu') if data_path is not None: dataset = ImageDataset(data_path,lab=lab,pretrafo=trafo) loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0) if input_image is not None: img=trafo(Image.open(input_image)) if lab: img=color.rgb2lab(np.asarray(img)/255)[...,:1]-np.array([50])[None,None,:] loader = [(transforms.ToTensor()(img)[None,...].float(),input_image)] classes=(150 if classification else 2) if lab else 3 #define model UNet=None zoom=False if mode == 0: UNet=model(col_channels=classes) elif mode ==1: UNet=unet(drop_rate=drop_rate,classes=classes) elif mode ==2: UNet=generator(drop_rate,classes) elif mode == 3: UNet=richzhang(drop_rate,classes) zoom=True elif mode == 4: UNet=color_unet(True,drop_rate,classes) elif mode == 5: UNet = middle_unet(True,drop_rate,classes) #load weights try: UNet.load_state_dict(torch.load(weight_path, map_location=device)) print("Loaded network weights from", weight_path) except FileNotFoundError: print("Did not find weight files.") sys.exit() outpath=None UNet.to(device) UNet.eval() with torch.no_grad(): for i,(X,name) in enumerate(loader): X=X.to(device) unet_col=UNet(X) col=show_colorization(unet_col,original=X,lab=lab,cl=classification,zoom=zoom,T=temp,return_img=output is not None) if output: try: fp,f=os.path.split(name) except TypeError: fp,f=os.path.split(name[0]) n,e=f.split('.') f='.'.join((n+'_color','png')) outpath=output if os.path.isdir(output) or os.path.isdir(os.path.basename(output)) else fp Image.fromarray(toInt(col[0])).save(os.path.join(outpath,f)) if output: print('Finished colorization. Go to "%s" to see the colorized version(s) of the image(s)'%os.path.realpath(outpath))
def test(FLAGS): ### Parameters of the testing #cropShape = (192,192,3) cropShape = (64, 64, 3) test_path = FLAGS.test_path #noise_level = FLAGS.noise architecture = FLAGS.architecture filepath = FLAGS.load_weights_path sliding_window = FLAGS.sliding_window # Load model if architecture == "unet": model = unet(cropShape) elif architecture == "wavelet": model = unetWavelet(cropShape) else: raise RuntimeError( 'Unkwown architecture, please choose from "unet" or "wavelet".') model.load_weights(filepath) # Load the images in the test folder test_files = [img_path for img_path in glob(test_path + '/*.png')] i = 1 for img_path in test_files: image = load_img(img_path) image = img_to_array(image) if sliding_window != 0: image = image[:sliding_window * ((image.shape[0] - cropShape[0]) // sliding_window) + cropShape[0], :sliding_window * ((image.shape[1] - cropShape[1]) // sliding_window) + cropShape[1]] else: image = image[:(image.shape[0] // cropShape[0]) * cropShape[0], :(image.shape[1] // cropShape[1]) * cropShape[1]] ## Add noise #noise = np.random.normal(0,noise_level,image.shape) #noisy_image = image + noise #noisy_image /= 255.0 #residual_image = np.empty(image.shape) blurred_image = add_motion_blur(image) blurred_image /= 255.0 residual_image = np.empty(image.shape) # Divide the image into subimages of size cropShape to match the model's input size. if sliding_window != 0: sum_img = np.zeros(image.shape) count_img = np.zeros(image.shape) for y in range(0, sliding_window + image.shape[0] - cropShape[0], sliding_window): for x in range(0, sliding_window + image.shape[1] - cropShape[1], sliding_window): extract = blurred_image[y:y + cropShape[0], x:x + cropShape[1]] sum_img[y:y + cropShape[0], x:x + cropShape[1]] += model.predict(extract[np.newaxis, ...])[0] count_img[y:y + cropShape[0], x:x + cropShape[1]] += 1 residual_image = sum_img / count_img else: for block_y in range(0, image.shape[0] // cropShape[0]): for block_x in range(0, image.shape[1] // cropShape[1]): extract = blurred_image[block_y * cropShape[0]:(block_y + 1) * cropShape[0], block_x * cropShape[1]:(block_x + 1) * cropShape[1]] residual_image[block_y * cropShape[0]:(block_y + 1) * cropShape[0], block_x * cropShape[1]:(block_x + 1) * cropShape[1]] = model.predict( extract[np.newaxis, ...])[0] #normalize_img = np.zeros(noisy_image.shape) #normalize_img = # Display the noised and denoised images #display_img(noisy_image,residual_image) cv.imwrite('wavelet_rot05_blur_20/blurred/' + str(i) + '.png', blurred_image * 255) cv.imwrite('wavelet_rot05_blur_20/reconstructed/' + str(i) + '.png', (blurred_image - residual_image) * 255) cv.imwrite('wavelet_rot05_blur_20/original/' + str(i) + '.png', image) #np.save('unet_noisy/original/'+str(i)+'.npy',noisy_image) #np.save('unet_noisy/reconstructed/'+str(i)+'.npy',residual_image) i += 1
model = unet_dilated(input_size=input_shape) elif is_imageNet: model_unet = Unet(BACKBONE, encoder_weights='imagenet') if is_stacked: new_model = keras.models.Sequential() new_model.add( Conv2D(3, (1, 1), padding='same', activation='relu', input_shape=input_shape)) new_model.add(model_unet) model = new_model else: model = model_unet else: model = unet(input_size=input_shape) model.compile(loss='binary_crossentropy', optimizer=Adam(lr=learning_rate_scheduler(0)), metrics=['acc', f1]) checkpoint = ModelCheckpoint(filepath, monitor='f1', verbose=1, save_best_only=True, mode='max') csv_logger = CSVLogger(csv_log_file, append=True, separator=';') callbacks_list = [checkpoint, csv_logger] model.fit_generator(batch_generator(train_df, batch_size, is_stacked, is_imageNet),
def _unet_model_fn(features, labels, mode, params=None, config=None, model_dir=None): features = tf.reshape( features, [params['batch_size'], params['resolution'], params['resolution'], 1]) training = (mode == tf.estimator.ModeKeys.TRAIN) result = unet(features, 1, params['num_chans'], params['drop_prob'], params['num_pools'], training) loss = None train_op = None hooks = [] export_outputs = None eval_hooks = [] chief_hooks = [] metrics = {} if mode != tf.estimator.ModeKeys.PREDICT: learning_rate_var = tf.Variable( float(params['lr']), trainable=False, name='lr', collections=[tf.GraphKeys.LOCAL_VARIABLES]) loss = tf.losses.absolute_difference(labels, result) mse = tf.losses.mean_squared_error(labels, result) nmse = tf.norm(labels - result)**2 / tf.norm(labels)**2 global_step = tf.train.get_or_create_global_step() epoch = global_step // params['epoch_len'] if training: tf.summary.scalar('lr', learning_rate_var) tf.summary.scalar('mse', mse) tf.summary.scalar('nmse', nmse) board_hook = MlBoardReporter( { "_step": global_step, "_epoch": epoch, "_train_loss": loss, '_train_lr': learning_rate_var, '_train_mse': mse, '_train_nmse': nmse }, every_steps=params['save_summary_steps']) chief_hooks = [board_hook] update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): opt = tf.train.RMSPropOptimizer(learning_rate_var, params['weight_decay']) train_op = opt.minimize(loss, global_step=global_step) rimage = (result - tf.reduce_min(result)) rimage = rimage / tf.reduce_max(rimage) tf.summary.image('Reconstruction', rimage, 3) limage = (labels - tf.reduce_min(labels)) limage = limage / tf.reduce_max(limage) tf.summary.image('Original', limage, 3) hooks = [ TrainingLearningRateHook(params['epoch_len'], learning_rate_var, float(params['lr']), int(params['lr_step_size']), float(params['lr_gamma'])) ] if not training: metrics['mse'] = tf.metrics.mean(mse) metrics['nmse'] = tf.metrics.mean(nmse) eval_hooks = [ tf.train.SummarySaverHook(save_steps=1, output_dir=model_dir + "/test", summary_op=tf.summary.merge_all()) ] else: export_outputs = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: tf.estimator.export.PredictOutput(result) } return tf.estimator.EstimatorSpec(mode=mode, eval_metric_ops=metrics, predictions=result, training_chief_hooks=chief_hooks, loss=loss, training_hooks=hooks, export_outputs=export_outputs, evaluation_hooks=eval_hooks, train_op=train_op)
import os import tensorflow as tf from models.unet import unet import keras as keras parser = argparse.ArgumentParser( description='convert unet Model with h5 encoded weights to pb') parser.add_argument("--weights", help="path to h5 encoded weights") args = parser.parse_args() keras.backend.clear_session() keras.backend.set_learning_phase(0) with keras.backend.get_session() as sess: model = unet(input_shape=(256, 512, 3), num_classes=4, lr_init=1e-3, lr_decay=5e-4) model.load_weights(args.weights) # save frozen subset graph for acceleration output_graph = os.path.splitext(args.weights)[0] #+ ".pb" input_names = [out.op.name for out in model.inputs] print(input_names) output_names = [out.op.name for out in model.outputs] print('input node is{}'.format(input_names)) print('output node is{}'.format(output_names)) saver = tf.train.Saver() graph_def = sess.graph.as_graph_def() # for node in graph_def.node:
def main(argv): data_path = s.data_path weight_path = s.weights_path mode = 1 drop_rate = 0 lab = s.lab classification = False temp = .4 try: opts, args = getopt.getopt(argv, "h:w:p:b:m:ld:ct:", [ "help", "weight-path=", "datapath=", 'model=', 'lab', 'drop-rate=' ]) except getopt.GetoptError as error: print(error) print( 'test.py -w <path to weights file> -p <path to dataset> -l <no argument. use if lab should be used> -m <mode: different models>\ -d <amount of dropout used in model> -c <no argument. Use if model is classifier> -t <temperature for annealed mean> -o <output path for images>' ) sys.exit(2) print("opts", opts) for opt, arg in opts: if opt == '-h': print('test.py -i <Boolean> -s <Boolean>') sys.exit() elif opt in ("-w", "--weight-path"): weight_path = arg elif opt in ("--datapath", "-p"): data_path = arg elif opt in ("--batchnorm", "-b"): batch_norm = arg in ["True", "true", "1"] elif opt == '-m': if arg in ('custom', '0'): mode = 0 elif arg in ('u', '1', 'unet'): mode = 1 elif arg in ('ende', '2'): mode = 2 elif arg in ('richzhang', 'classende', '3'): mode = 3 elif arg in ('colorunet', 'cu', '4'): mode = 4 elif arg in ('mu', '5', 'middle'): mode = 5 elif opt in ('-l', '--lab'): lab = True elif opt in ("-d", "--drop-rate"): drop_rate = float(arg) elif opt == '-c': classification = True lab = True elif opt == '-t': temp = float(arg) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') dataset = None if data_path == './cifar-10': in_size = 32 dataset = 0 elif 'places' in data_path: in_size = 224 dataset = 1 elif 'stl' in data_path: in_size = 96 dataset = 2 in_shape = (3, in_size, in_size) #out_shape=(s.classes,32,32) trainset = load_trainset(data_path, train=False, lab=lab) trainloader = torch.utils.data.DataLoader( trainset, batch_size=3, shuffle=True, num_workers=2 if dataset in (0, 1) else 0) print("Loaded dataset from", data_path) classes = (150 if classification else 2) if lab else 3 #define model UNet = None zoom = False if mode == 0: UNet = model(col_channels=classes) elif mode == 1: UNet = unet(drop_rate=drop_rate, classes=classes) elif mode == 2: UNet = generator(drop_rate, classes) elif mode == 3: UNet = richzhang(drop_rate, classes) zoom = True elif mode == 4: UNet = color_unet(True, drop_rate, classes) elif mode == 5: UNet = middle_unet(True, drop_rate, classes) #load weights try: UNet.load_state_dict(torch.load(weight_path, map_location=device)) print("Loaded network weights from", weight_path) except FileNotFoundError: print("Did not find weight files.") #sys.exit(2) UNet.to(device) UNet.eval() gray = torch.tensor([0.2989, 0.5870, 0.1140])[:, None, None].float() with torch.no_grad(): for i, batch in enumerate(trainloader): if dataset == 0: #cifar 10 (image, _) = batch elif dataset in (1, 2): #places image = batch X = None if lab: if dataset == 0: #cifar 10 image = np.transpose(image, (0, 2, 3, 1)) image = np.transpose(color.rgb2lab(image), (0, 3, 1, 2)) image = torch.from_numpy( (image - np.array([50, 0, 0])[None, :, None, None])).float() X = torch.unsqueeze(image[:, 0, :, :], 1).to( device) #set X to the Lightness of the image image = image[:, 1:, :, :] #image is a and b channel else: #convert to grayscale image #using the matlab formula: 0.2989 * R + 0.5870 * G + 0.1140 * B and load data to gpu X = (image.clone() * gray).sum(1).to(device).view( -1, 1, *in_shape[1:]) image = image.float() #print(X.min(),X.max()) #generate colorized version with unet #for arr in (image[:,0,...],image[:,1,...],X): # print(arr.min(),arr.max()) try: unet_col = UNet(X) except: unet_col = UNet(torch.stack((X, X, X), 1)[:, :, 0, :, :]) #for arr in (unet_col[0,...],unet_col[1,...]): # print(arr.min().item(),arr.max().item()) show_colorization(unet_col, image, X, lab=lab, cl=classification, zoom=zoom, T=temp)
def model_fn(features, labels, mode, params=None, config=None, model_dir=None): embedding_dim = 128 if mode == tf.estimator.ModeKeys.PREDICT: x = features['images'] else: x = features x = unet(x, embedding_dim, 32, 0.5, 4, mode == tf.estimator.ModeKeys.TRAIN) logging.info('Unet {}'.format(x.shape)) word_index = params['word_index'] x = tf.nn.relu(x) _, l1, l2, _ = tf.unstack(tf.shape(x)) x = tf.reshape(x, [params['batch_size'], l1 * l2, embedding_dim]) features_length = tf.zeros( (params['batch_size']), dtype=tf.int64) + tf.cast(l1 * l2, tf.int64) with tf.variable_scope('embedding'): embedding = tf.get_variable('embedding_op', [len(params['word_index']), embedding_dim], initializer=tf.random_uniform_initializer( -1, 1), trainable=True, dtype=tf.float32) start_tokens = tf.zeros([params['batch_size']], dtype=tf.int32) + word_index['<start>'] if mode == tf.estimator.ModeKeys.PREDICT: helper = tf.contrib.seq2seq.GreedyEmbeddingHelper( embedding, start_tokens=tf.to_int32(start_tokens), end_token=word_index['<end>']) else: train_output = tf.concat([tf.expand_dims(start_tokens, 1), labels], 1) output_lengths = tf.reduce_sum( tf.to_int32(tf.not_equal(train_output, word_index['<end>'])), 1) output_embed = tf.nn.embedding_lookup(embedding, train_output) if mode == tf.estimator.ModeKeys.EVAL: helper = tf.contrib.seq2seq.TrainingHelper(output_embed, output_lengths) else: helper = tf.contrib.seq2seq.TrainingHelper(output_embed, output_lengths) #helper = tf.contrib.seq2seq.ScheduledEmbeddingTrainingHelper(output_embed, output_lengths, embedding, 0.5) cells = [ tf.contrib.rnn.GRUCell(params['hidden_size']) for _ in range(params['num_layers']) ] mrnn = tf.contrib.rnn.MultiRNNCell(cells, state_is_tuple=True) attention_mechanism = tf.contrib.seq2seq.BahdanauAttention( num_units=params['hidden_size'], memory=x, memory_sequence_length=features_length) attn_cell = tf.contrib.seq2seq.AttentionWrapper( mrnn, attention_mechanism, attention_layer_size=params['hidden_size'], alignment_history=(mode == tf.estimator.ModeKeys.PREDICT)) output_layer = tf.layers.Dense(len(word_index)) initial_state = attn_cell.zero_state(dtype=tf.float32, batch_size=params['batch_size']) decoder = tf.contrib.seq2seq.BasicDecoder(attn_cell, helper, initial_state, output_layer=output_layer) outputs = tf.contrib.seq2seq.dynamic_decode( decoder, output_time_major=False, impute_finished=True, maximum_iterations=params['max_target_seq_length']) train_op = None if mode == tf.estimator.ModeKeys.PREDICT: loss = None predictions = outputs[0].sample_id alligments = outputs[1].alignment_history.stack() export_outputs = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: tf.estimator.export.PredictOutput({ 'labels': predictions, 'attentions': alligments }) } else: predictions = None export_outputs = None train_outputs = outputs[0] logging.info('Output: {}'.format(train_outputs.rnn_output)) weights = tf.to_float( tf.not_equal(train_output[:, :-1], word_index['<end>'])) loss = tf.contrib.seq2seq.sequence_loss(train_outputs.rnn_output, labels, weights=weights) if mode == tf.estimator.ModeKeys.TRAIN: opt = tf.train.AdamOptimizer(params['learning_rate']) if params['grad_clip'] is None: train_op = opt.minimize( loss, global_step=tf.train.get_or_create_global_step()) else: gradients, variables = zip(*opt.compute_gradients(loss)) gradients, _ = tf.clip_by_global_norm(gradients, params['grad_clip']) train_op = opt.apply_gradients( [(gradients[i], v) for i, v in enumerate(variables)], global_step=tf.train.get_or_create_global_step()) return tf.estimator.EstimatorSpec(mode=mode, eval_metric_ops=None, predictions=predictions, loss=loss, export_outputs=export_outputs, train_op=train_op)
def train(FLAGS): ### Parameters of the training # Can be freely changed epochs = 400 train_ratio = 0.8 batch_size = 16 cropShape = (192,192,3) early_stop = 10 # The learning rate is decayed from 10e-4 to 10e-5 over 200 epochs. optimizer = Adam(lr=10e-4, beta_1=0.9, beta_2=0.999, epsilon=10e-8, decay=10e-3) train_path = FLAGS.train_path valid_path = FLAGS.test_path noise_level = FLAGS.noise architecture = FLAGS.architecture if architecture == "unet": filepath = "weights/DenoisingUnet"+"_"+str(noise_level)+".h5" model = unet(cropShape) elif architecture == "wavelet": filepath = "weights/DenoisingWavelet"+"_"+str(noise_level)+".h5" model = unetWavelet(cropShape) else : raise RuntimeError('Unkwown architecture, please choose from "unet" or "wavelet".') #model.summary() model.compile(loss='mean_squared_error', optimizer=optimizer, metrics=[psnr]) ### Preparation of the dataset, building the generators train_files =[img_path for img_path in glob(train_path + '/*.png')] test_files =[img_path for img_path in glob(valid_path + '/*.png')] train_generator = CustomGenerator(train_files[0:int(len(train_files)*train_ratio)], noise_level, batch_size, cropShape) valid_generator = CustomGenerator(train_files[int(len(train_files)*train_ratio):], noise_level, batch_size, cropShape) test_generator = CustomGenerator(test_files, noise_level, batch_size, cropShape) ### Train the model model.fit_generator( train_generator, epochs=epochs, verbose=2, callbacks = [ ModelCheckpoint("temp.h5", monitor='val_loss', verbose=0, save_best_only=True, mode='auto'), EarlyStopping(monitor='val_loss', patience=early_stop, verbose=0, mode='auto') ], validation_data=valid_generator ) model.load_weights("temp.h5") model.save(filepath) print("Training done") ### Evaluate the model on the test dataset result = model.evaluate_generator(test_generator) print(result)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' if config.model_settings.no_cuda: os.environ['CUDA_VISIBLE_DEVICES'] = '-1' else: if not len(tf.config.list_physical_devices("GPU")): print("Could not find GPU, exiting. Change config.yaml\ to turn off GPU") exit(1) sf1 = StreamingF1Score( num_classes=config.model_settings.num_classes, focus_on_class=config.model_settings.f1_focus_on_class) model = unet.unet( (None, None, config.model_settings.unet_input_channels), n_classes=config.model_settings.num_classes, initial_exp=config.model_settings.unet_initial_exp, weight_decay_const=config.model_settings.weight_decay_const, apply_batchnorm=config.model_settings.apply_batchnorm) lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay( initial_learning_rate=config.model_settings.initial_learning_rate, decay_steps=config.model_settings.decay_steps, decay_rate=config.model_settings.decay_rate, staircase=config.model_settings.staircase) optim = tf.keras.optimizers.Adam(learning_rate=lr_schedule) model.compile(optim, loss=config.model_settings.objective, metrics=[m_acc, sf1])
args = parser.parse_args() if __name__ == '__main__': models_root = '/data/dev/models' cudev = args.gpu_num transformations = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # set model os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_num) device, gpus = utils.set_device() checkpoint = torch.load(args.path_to_model) model = unet.unet(3, pretrained=False) model.load_state_dict(checkpoint['state_dict']) model.to(device) model.eval() torch.set_grad_enabled(False) # find paths vid_paths = [] with open(args.videos) as a: for line in a: path = line.rstrip() vid_paths.append(path) frames_list = [] masks_list = [] bboxes_list = [] landmarks_list = []
# files and paths TRAIN_DIR = results.SRC_DIR for d in [WEIGHT_DIR, TB_LOG_DIR]: if not os.path.exists(d): os.makedirs(d) PATCH_SIZE = [int(x) for x in results.patch_size.split("x")] ######### MODEL AND CALLBACKS ######### if not model: model = unet( model_path=MODEL_PATH, num_channels=num_channels, loss=continuous_dice_coef_loss, ds=2, lr=learning_rate, num_gpus=NUM_GPUS, verbose=1, ) else: print("Continuing training with", model) model = load_model(model, custom_objects=custom_losses) monitor = "val_dice_coef" # checkpoints checkpoint_filename = str(start_time) +\ "_epoch_{epoch:04d}_" +\ monitor+"_{"+monitor+":.4f}_weights.hdf5"
def main(argv): # setting argument defaults mbsize = s.batch_size report_freq = s.report_freq weight_path = s.weights_path weights_name = s.weights_name lr = s.learning_rate save_freq = s.save_freq mode = 3 epochs = s.epochs beta1, beta2 = s.betas infinite_loop = s.infinite_loop data_path = s.data_path drop_rate = 0 lab = s.lab load_list = s.load_list help = 'train_regression.py -b <batch size> -e <amount of epochs to train. standard: infinite> -r <report frequency> -w <path to weights folder> \ -n <name> -s <save freq.> -l <learning rate> -p <path to data set> -d <dropout rate> -m <mode: differnet models> --beta1 <beta1 for adam>\ --beta2 <beta2 for adam> --lab <No argument. If used lab colorspace is used> \ --lambda <hyperparameter for class weights>' try: opts, args = getopt.getopt(argv, "he:b:r:w:l:s:n:m:p:d:i:", [ 'epochs=', "mbsize=", "report-freq=", 'weight-path=', 'lr=', 'save-freq=', 'weight-name=', 'mode=', 'data_path=', 'drop_rate=' 'beta1=', 'beta2=', 'lab', 'image-loss-weight=', 'load-list' ]) except getopt.GetoptError: print(help) sys.exit(2) print("opts", opts) for opt, arg in opts: if opt == '-h': print(help) sys.exit() elif opt in ("-b", "--mbsize"): mbsize = int(arg) elif opt in ("-e", "--epochs"): epochs = int(arg) infinite_loop = False elif opt in ('-r', '--report-freq'): report_freq = int(arg) elif opt in ("-w", "--weight-path"): weight_path = arg elif opt in ("-n", "--weight-name"): weights_name = arg elif opt in ("-s", "--save-freq"): save_freq = int(arg) elif opt in ("-l", "--lr"): lr = float(arg) elif opt == '-m': if arg in ('custom', '0'): mode = 0 elif arg in ('u', '1', 'unet'): mode = 1 elif arg in ('ende', '2'): mode = 2 elif arg in ('color', '3', 'cu'): mode = 3 elif opt in ("-p", "--data_path"): data_path = str(arg) elif opt in ("-d", "--drop_rate"): drop_rate = float(arg) elif opt == '--beta1': beta1 = float(arg) elif opt == '--beta2': beta2 = float(arg) elif opt == '--lab': lab = True elif opt in ('--load-list'): load_list = True device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') dataset = None if 'cifar' in data_path: in_size = 32 dataset = 0 elif 'places' in data_path: in_size = 224 dataset = 1 elif 'stl' in data_path: in_size = 96 dataset = 2 in_shape = (3, in_size, in_size) #out_shape=(s.classes,32,32) betas = (beta1, beta2) weight_path_ending = os.path.join(weight_path, weights_name + '.pth') loss_path_ending = os.path.join(weight_path, weights_name + "_" + s.loss_name) trainset = load_trainset(data_path, lab=lab, load_list=load_list) trainloader = torch.utils.data.DataLoader(trainset, batch_size=mbsize, shuffle=True, num_workers=2) print("NETWORK PATH:", weight_path_ending) #define output channels of the model classes = 2 if lab else 3 #define model UNet = None if mode == 0: UNet = model(col_channels=classes) elif mode == 1: UNet = unet(drop_rate=drop_rate, classes=classes) elif mode == 2: UNet = generator(drop_rate, classes) elif mode == 3: UNet = color_unet(True, drop_rate, classes) #load weights try: UNet.load_state_dict( torch.load(weight_path_ending, map_location=device)) print("Loaded network weights from", weight_path) except FileNotFoundError: print("Initialize new weights for the generator.") UNet.to(device) #save the hyperparameters to a JSON-file for better oranization model_description_path_ending = os.path.join(weight_path, s.model_description_name) # initialize model dict try: with open(model_description_path_ending, "r") as file: model_dict = json.load(file) except FileNotFoundError: model_dict = {} prev_epochs = 0 # save settings in dict if new weights are beeing initialized if not weights_name in model_dict.keys(): model_dict[weights_name] = { "loss_name": loss_path_ending, "epochs": 0, "batch_size": mbsize, "lr": lr, "lab": lab, "betas": betas, "model": ['custom', 'unet', 'encoder-decoder', 'color-unet'][mode] } else: #load specified parameters from model_dict params = model_dict[weights_name] mbsize = params['batch_size'] betas = params['betas'] lr = params['lr'] lab = params['lab'] loss_path_ending = params['loss_name'] #memorize how many epochs already were trained if we continue training prev_epochs = params['epochs'] + 1 #optimizer optimizer_g = optim.Adam(UNet.parameters(), lr=lr, betas=betas) # l1 loss l1loss = nn.L1Loss().to(device) loss_hist = [] UNet.train() gray = torch.tensor([0.2989, 0.5870, 0.1140])[:, None, None].float() # run over epochs for e in (range(prev_epochs, prev_epochs + epochs) if not infinite_loop else count(prev_epochs)): g_running = 0 #load batches for i, batch in enumerate(trainloader): if dataset == 0: #cifar 10 (image, _) = batch elif dataset in (1, 2): #places and stl 10 image = batch X = None #differentiate between the two available color spaces RGB and Lab if lab: if dataset == 0: #cifar 10 image = np.transpose(image, (0, 2, 3, 1)) image = np.transpose(color.rgb2lab(image), (0, 3, 1, 2)) image = torch.from_numpy( (image + np.array([-50, 0, 0])[None, :, None, None])).float() X = torch.unsqueeze(image[:, 0, :, :], 1).to( device) #set X to the Lightness of the image image = image[:, 1:, :, :].to(device) #image is a and b channel else: #convert to grayscale image #using the matlab formula: 0.2989 * R + 0.5870 * G + 0.1140 * B and load data to gpu X = (image.clone() * gray).sum(1).to(device).view( -1, 1, *in_shape[1:]) image = image.float().to(device) #---------------------------------------------------------------------------------------- ################################### Unet optimization ################################### #---------------------------------------------------------------------------------------- #clear gradients optimizer_g.zero_grad() #generate colorized version with unet unet_col = None #print(X.shape,image.shape,classes) if mode == 0: unet_col = UNet(torch.stack((X, X, X), 1)[:, :, 0, :, :]) else: unet_col = UNet(X) #calculate how close the generated pictures are to the ground truth loss_g = l1loss(unet_col, image) #backpropagation loss_g.backward() optimizer_g.step() g_running += loss_g.item() loss_hist.append([e, i, loss_g.item()]) #report running loss if (i + len(trainloader) * e) % report_freq == report_freq - 1: print('Epoch %i, batch %i: \tunet loss=%.2e' % (e + 1, i + 1, g_running / report_freq)) g_running = 0 if s.save_weights and ( i + len(trainloader) * e) % save_freq == save_freq - 1: #save parameters try: torch.save(UNet.state_dict(), weight_path_ending) except FileNotFoundError: os.makedirs(weight_path) torch.save(UNet.state_dict(), weight_path_ending) print("Parameters saved") if s.save_loss: #save loss history to file try: f = open(loss_path_ending, 'a') np.savetxt(f, loss_hist, '%e') f.close() except FileNotFoundError: os.makedirs(s.loss_path) np.savetxt(loss_path_ending, loss_hist, '%e') loss_hist = [] #update epoch count in dict after each epoch model_dict[weights_name]["epochs"] = e #save it to file try: with open(model_description_path_ending, "w") as file: json.dump(model_dict, file, sort_keys=True, indent=4) except: print('Could not save to model dictionary (JSON-file)')
def main(argv): # setting argument defaults mbsize = s.batch_size report_freq = s.report_freq weight_path = s.weights_path weights_name = s.weights_name lr = s.learning_rate save_freq = s.save_freq mode = 3 image_loss_weight = s.image_loss_weight epochs = s.epochs beta1, beta2 = s.betas infinite_loop = s.infinite_loop data_path = s.data_path drop_rate = 0 lab = True weighted_loss = True weight_lambda = .25 load_list = s.load_list help = 'train_classification.py -b <batch size> -e <amount of epochs to train. standard: infinite> -r <report frequency> -w <path to weights folder> \ -n <name> -s <save freq.> -l <learning rate> -p <path to data set> -d <dropout rate> -m <mode: differnet models> --beta1 <beta1 for adam>\ --beta2 <beta2 for adam> --lab <No argument. If used lab colorspace is cused> --weighted <No argument. If used *NO* class weights are used> \ --lambda <hyperparameter for class weights>' try: opts, args = getopt.getopt(argv, "he:b:r:w:l:s:n:p:d:i:m:", [ 'epochs=', "mbsize=", "report-freq=", 'weight-path=', 'lr=', 'save-freq=', 'weight-name=', 'data_path=', 'drop_rate=' 'beta1=', 'beta2=', 'lab', 'image-loss-weight=', 'weighted', 'mode=', 'lambda=' ]) except getopt.GetoptError: print(help) sys.exit(2) print("opts", opts) for opt, arg in opts: if opt == '-h': print(help) sys.exit() elif opt in ("-b", "--mbsize"): mbsize = int(arg) elif opt in ("-e", "--epochs"): epochs = int(arg) infinite_loop = False elif opt in ('-r', '--report-freq'): report_freq = int(arg) elif opt in ("-w", "--weight-path"): weight_path = arg elif opt in ("-n", "--weight-name"): weights_name = arg elif opt in ("-s", "--save-freq"): save_freq = int(arg) elif opt in ("-l", "--lr"): lr = float(arg) elif opt in ("-p", "--data_path"): data_path = str(arg) elif opt in ("-d", "--drop_rate"): drop_rate = float(arg) elif opt == '-m': if arg in ('richzhang', '0', 'ende'): mode = 0 elif arg in ('u', '1', 'unet'): mode = 1 elif arg in ('color', '2', 'cu'): mode = 2 elif arg in ('mu', '3', 'middle'): mode = 3 elif opt == '--beta1': beta1 = float(arg) elif opt == '--beta2': beta2 = float(arg) elif opt == '--lab': lab = True elif opt == '--weighted': weighted_loss = not weighted_loss elif opt == '--load-list': load_list = not load_list elif opt == '--lambda': weight_lambda = float(arg) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') dataset = None in_size = 256 if 'cifar' in data_path: in_size = 32 dataset = 0 elif 'places' in data_path: in_size = 224 dataset = 1 elif 'stl' in data_path: in_size = 96 dataset = 2 in_shape = (3, in_size, in_size) #out_shape=(s.classes,32,32) betas = (beta1, beta2) weight_path_ending = os.path.join(weight_path, weights_name + '.pth') loss_path_ending = os.path.join(weight_path, weights_name + "_" + s.loss_name) trainset = load_trainset(data_path, lab=lab, normalize=False, load_list=load_list) trainloader = torch.utils.data.DataLoader( trainset, batch_size=mbsize, shuffle=True, num_workers=2 if dataset in (0, 1) else 0) print("NETWORK PATH:", weight_path_ending) #define output channels of the model classes = 150 #define model if mode == 0: classifier = generator(drop_rate, classes) elif mode == 1: classifier = unet(True, drop_rate, classes) elif mode == 2: classifier = color_unet(True, drop_rate, classes) elif mode == 3: classifier = middle_unet(True, drop_rate, classes) #load weights try: classifier.load_state_dict(torch.load(weight_path_ending)) print("Loaded network weights from", weight_path) except FileNotFoundError: print("Initialize new weights for the generator.") #sys.exit(2) classifier.to(device) #save the hyperparameters to a JSON-file for better oranization model_description_path_ending = os.path.join(weight_path, s.model_description_name) # initialize model dict try: with open(model_description_path_ending, "r") as file: model_dict = json.load(file) except FileNotFoundError: model_dict = {} prev_epochs = 0 # save settings in dict if new weights are beeing initialized if not weights_name in model_dict.keys(): model_dict[weights_name] = { "loss_name": loss_path_ending, "epochs": 0, "batch_size": mbsize, "lr": lr, "lab": lab, "betas": betas, "image_loss_weight": image_loss_weight, "weighted_loss": weighted_loss, "model": 'classification ' + ['richzhang', 'U-Net', 'color U-Net', 'middle U-Net'][mode] } else: #load specified parameters from model_dict params = model_dict[weights_name] #mbsize=params['batch_size'] betas = params['betas'] #lr=params['lr'] lab = params['lab'] image_loss_weight = params['image_loss_weight'] weighted_loss = params['weighted_loss'] loss_path_ending = params['loss_name'] #memorize how many epochs already were trained if we continue training prev_epochs = params['epochs'] + 1 #optimizer optimizer = optim.Adam(classifier.parameters(), lr=lr, betas=betas) class_weight_path = 'resources/class-weights.npy' if weighted_loss: weights = np.load(class_weight_path) if dataset == 0: class_weight_path = 'resources/cifar-lab-class-weights.pt' weights = torch.load(class_weight_path).numpy() elif dataset == 2: if weight_lambda: class_weight_path = 'resources/probdist_lab.pt' prob_dict = torch.load(class_weight_path) prob = np.array(list(prob_dict.values())) weights = 1 / ((1 - weight_lambda) * prob / prob.sum() + weight_lambda / classes) else: class_weight_path = 'resources/class-weights-lab150-stl.pt' weights = torch.load(class_weight_path) print('Class-weights loaded from ' + class_weight_path) criterion = softCossEntropyLoss( weights=weights, device=device) if weighted_loss else softCossEntropyLoss(weights=None, device=device) loss_hist = [] soft_onehot = torch.load('resources/smooth_onehot150.pt', map_location=device) classifier.train() # run over epochs for e in (range(prev_epochs, prev_epochs + epochs) if not infinite_loop else count(prev_epochs)): g_running = 0 #load batches for i, batch in enumerate(trainloader): if dataset == 0: #cifar 10 (image, _) = batch elif dataset in (1, 2): #places image = batch #batch_size=image.shape[0] if dataset == 0: #cifar/stl 10 image = np.transpose(image, (0, 2, 3, 1)) image = np.transpose(color.rgb2lab(image), (0, 3, 1, 2)) image = torch.from_numpy( (image - np.array([50, 0, 0])[None, :, None, None])).float() X = image[:, :1, :, :].to( device) #set X to the Lightness of the image image = image[:, 1:, :, :].to(device) #image is a and b channel #---------------------------------------------------------------------------------------- ################################### Model optimization ################################## #---------------------------------------------------------------------------------------- #clear gradients optimizer.zero_grad() #softmax activated distribution model_out = classifier(X).double() #create bin coded verion of ab ground truth binab = ab2bins(image.transpose(1, 3).transpose(1, 2)) if mode == 0: binab = F.interpolate(binab.float(), scale_factor=(.25, .25)).long() binab = torch.squeeze(binab, 1) binab = soft_onehot[:, binab].transpose(0, 1).double() #calculate loss loss = criterion(model_out, binab).mean(0) loss.backward() optimizer.step() g_running += loss.item() loss_hist.append([e, loss.item()]) #report running loss if (i + len(trainloader) * e) % report_freq == report_freq - 1: print('Epoch %i, batch %i: \tloss=%.2e' % (e + 1, i + 1, g_running / report_freq)) g_running = 0 if s.save_weights and ( i + len(trainloader) * e) % save_freq == save_freq - 1: #save parameters try: torch.save(classifier.state_dict(), weight_path_ending) #torch.save(crit.state_dict(),crit_path) except FileNotFoundError: os.makedirs(weight_path) torch.save(classifier.state_dict(), weight_path_ending) #torch.save(crit.state_dict(),crit_path) print("Parameters saved") if s.save_loss: #save loss history to file try: f = open(loss_path_ending, 'a') np.savetxt(f, loss_hist, '%e') f.close() except FileNotFoundError: os.makedirs(s.loss_path) np.savetxt(loss_path_ending, loss_hist, '%e') loss_hist = [] #update epoch count in dict after each epoch model_dict[weights_name]["epochs"] = e #save it to file try: with open(model_description_path_ending, "w") as file: json.dump(model_dict, file, sort_keys=True, indent=4) except: print('Could not save to model dictionary (JSON-file)')
def main(argv): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') gray = torch.tensor([0.2989 ,0.5870, 0.1140])[:,None,None].float() alex = alexnet().to(device) alex.load_state_dict(torch.load('weights/alexnet.pth')) alex.eval() classifier_path = 'weights/alexNorm.pt' weight_path = None mbsize = 16 col_space = None class_net = False mode = 2 drop_rate = 0 T = .4 classes = 3 stl_path = './stl-10' help='metric.py -b <batch size> -w <model weight file> -p <2layer classifer weight file> -s <colorspace "rgb" or "lab"> \ -c <use if model is classification model> -m <mode> -d <dropout rate> -t <define temperature> --stl <path to stl dataset> try: opts, args = getopt.getopt(argv,"b:w:p:cm:d:t:s:", ["mbsize=",'weight-path=','classifier-path=','c_space=','mode=','drop_rate=','temperature=','stl=']) except getopt.GetoptError: print(help) sys.exit(2) print("opts" ,opts) for opt, arg in opts: if opt in ("-b", "--mbsize"): mbsize = int(arg) elif opt in ("-w", "--weight-path"): weight_path = str(arg) elif opt in ("-p", "--classifier-path"): classifier_path = str(arg) elif opt in ("-s","--c_space"): col_space = str(arg) elif opt =='-c': class_net = True elif opt == '-m': if arg in ('u','0','unet'): mode = 0 elif arg in ('color','1','cu'): mode = 1 elif arg in ('mu','2','middle'): mode = 2 elif opt in ("-d", "--drop_rate"): drop_rate = float(arg) elif opt in ('-t', '--temperature'): T = float(arg) elif opt in ('--stl'): stl_path=arg if col_space == None: print('Specify color space') sys.exit(2) if class_net: if col_space == 'yuv': classes = 42 elif col_space == 'lab': classes = 150 if mode == 0: Col_Net = unet(True, drop_rate, classes).to(device) elif mode == 1: Col_Net = color_unet(True, drop_rate, classes).to(device) elif mode == 2: Col_Net = middle_unet(True, drop_rate, classes).to(device) Col_Net.load_state_dict(torch.load(weight_path, map_location=device)) Col_Net.eval() classifier=nn.Sequential(nn.Linear(1000,512),nn.ReLU(),nn.Linear(512,10)).to(device) classifier.load_state_dict(torch.load(classifier_path, map_location=device)) classifier.eval() #we trained on the test set and evaluate on the train set since the test set has more labeled images than the trainset testset = datasets.STL10(stl_path,split='train',transform=transforms.Compose([transforms.ToTensor()])) testloader = torch.utils.data.DataLoader(testset, batch_size=mbsize, shuffle=True, num_workers=0) print(pseudo_metric(testloader, col_space, Col_Net, classifier, alex, T))