def writeImage(result, file_name): file_name = utils.filepath_to_name(args.image) cv2.imwrite("%s_pred.png" % (file_name), cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) print("") print("Finished!") print("Wrote image " + "%s_pred.png" % (file_name))
def getCoordinatesFromFiles(images, outputPath=False, downscale_factor=False, background='white', saveMorphImage=True): coordinates = dict() for img in images: if not os.path.isfile(img): coordinates[utils.filepath_to_name(img, True)] = [] continue image = cv2.cvtColor(cv2.imread(img), cv2.COLOR_RGB2GRAY) if np.min(image) == np.max(image): #Whole image is same color i.e. everything is background coordinates[utils.filepath_to_name(img, True)] = [] continue if background == 'white': foregroundcolor = [np.min(image)] elif background == 'black': foregroundcolor = [np.max(image)] else: raise ValueError( 'Only black and white are allowed as values for background (it was {})' .format(background)) coordinates[utils.filepath_to_name( img, True)], morphImage = utils.getCoordsFromPrediction( image, foregroundcolor, downscale_factor, close_dim=384) #open_dim=8,,open2_dim=32 if saveMorphImage: outMorph = np.ones(morphImage.shape, dtype=np.uint8) * np.iinfo( np.uint8).max outMorph[morphImage == 0] = 0 cv2.imwrite(img[0:-4] + '_morph.tif', cv2.cvtColor(np.uint8(outMorph), cv2.COLOR_RGB2BGR)) if outputPath: if not outputPath.endswith('.json'): outputPath = outputPath + '.json' with open(outputPath, 'w') as f: json.dump(coordinates, f) return coordinates, images
sess.run(tf.global_variables_initializer()) print('Loading model checkpoint weights') saver = tf.train.Saver(max_to_keep=1000) saver.restore(sess, args.checkpoint_path) print("Testing image " + args.image) loaded_image = utils.load_image(args.image) resized_image = cv2.resize(loaded_image, (args.crop_width, args.crop_width)) input_image = np.expand_dims(np.float32( resized_image[:args.crop_height, :args.crop_width]), axis=0) / 255.0 st = time.time() output_image = sess.run(network, feed_dict={net_input: input_image}) run_time = time.time() - st output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation(output_image, label_values) file_name = utils.filepath_to_name(args.image) cv2.imwrite("%s_pred.png" % (file_name), cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) print("") print("Finished!") print("Wrote image " + "%s_pred.png" % (file_name))
gt = helpers.reverse_one_hot(helpers.one_hot_it(gt, label_values)) # st = time.time() output_image = sess.run(network, feed_dict={net_input: input_image}) output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation( output_image, label_values) accuracy, class_accuracies, prec, rec, f1, iou = utils.evaluate_segmentation( pred=output_image, label=gt, num_classes=num_classes) file_name = utils.filepath_to_name(val_input_names[ind]) target.write("%s, %f, %f, %f, %f, %f" % (file_name, accuracy, prec, rec, f1, iou)) for item in class_accuracies: target.write(", %f" % (item)) target.write("\n") scores_list.append(accuracy) class_scores_list.append(class_accuracies) precision_list.append(prec) recall_list.append(rec) f1_list.append(f1) iou_list.append(iou) gt = helpers.colour_code_segmentation(gt, label_values)
out_fore_clean = helpers.remove_small(out_fore_bin, min_size=300) out_fore_clean = helpers.remove_edge_seg(out_fore_clean) gt = helpers.remove_edge_seg(gt) out_vis_image = helpers.colour_code_segmentation(output_image, label_values) out_fore_vis_image = helpers.colour_code_segmentation(out_fore_clean, label_values) accuracy, class_accuracies, prec, rec, f1, iou = utils.evaluate_segmentation(pred=output_image, label=gt, num_classes=num_classes) accuracy_fore, class_accuracies_fore, prec_fore, rec_fore, f1_fore, iou_fore = utils.evaluate_segmentation(pred=out_fore_clean, label=gt, num_classes=num_classes) file_name = utils.filepath_to_name(test_input_names[ind]) target.write("%s, %f, %f, %f, %f, %f"%(file_name, accuracy, prec, rec, f1, iou)) for item in class_accuracies: target.write(", %f"%(item)) target.write("\n") scores_list.append(accuracy) class_scores_list.append(class_accuracies) precision_list.append(prec) recall_list.append(rec) f1_list.append(f1) iou_list.append(iou) scores_list_fore.append(accuracy_fore) class_scores_list_fore.append(class_accuracies_fore) precision_list_fore.append(prec_fore)
def main(args, _log): def data_augmentation(input_image, output_image): # Data augmentation input_image, output_image = utils.random_crop(input_image, output_image, args["crop_height"], args["crop_width"]) if args["h_flip"] and random.randint(0, 1): input_image = cv2.flip(input_image, 1) output_image = cv2.flip(output_image, 1) if args["v_flip"] and random.randint(0, 1): input_image = cv2.flip(input_image, 0) output_image = cv2.flip(output_image, 0) if args["brightness"]: factor = 1.0 + random.uniform(-1.0 * args["brightness"], args["brightness"]) table = np.array([((i / 255.0) * factor) * 255 for i in np.arange(0, 256)]).astype(np.uint8) input_image = cv2.LUT(input_image, table) if args["rotation"]: angle = random.uniform(-1 * args["rotation"], args["rotation"]) if args["rotation"]: M = cv2.getRotationMatrix2D( (input_image.shape[1] // 2, input_image.shape[0] // 2), angle, 1.0) input_image = cv2.warpAffine( input_image, M, (input_image.shape[1], input_image.shape[0]), flags=cv2.INTER_NEAREST) output_image = cv2.warpAffine( output_image, M, (output_image.shape[1], output_image.shape[0]), flags=cv2.INTER_NEAREST) return input_image, output_image print("Args:", args) # Get the names of the classes so we can record the evaluation results class_names_list, label_values = helpers.get_label_info( os.path.join(args["dataset"], "class_dict.csv")) class_names_string = "" for class_name in class_names_list: if not class_name == class_names_list[-1]: class_names_string = class_names_string + class_name + ", " else: class_names_string = class_names_string + class_name num_classes = len(label_values) config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) # Compute your softmax cross entropy loss net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3]) net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes]) network, init_fn = model_builder.build_model( model_name=args["model"], frontend=args["frontend"], net_input=net_input, num_classes=num_classes, crop_width=args["crop_width"], crop_height=args["crop_height"], is_training=True) loss = tf.reduce_mean( tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output)) opt = tf.train.RMSPropOptimizer( learning_rate=0.0001, decay=0.995).minimize( loss, var_list=[var for var in tf.trainable_variables()]) saver = tf.train.Saver(max_to_keep=1000) sess.run(tf.global_variables_initializer()) utils.count_params() # If a pre-trained ResNet is required, load the weights. # This must be done AFTER the variables are initialized with sess.run(tf.global_variables_initializer()) if init_fn is not None: init_fn(sess) # Load a previous checkpoint if desired model_checkpoint_name = "checkpoints/latest_model_" + args[ "model"] + "_" + basename(normpath(args["dataset"])) + ".ckpt" if args["continue_training"]: _log.info('Loaded latest model checkpoint') saver.restore(sess, model_checkpoint_name) # Load the data _log.info("Loading the data ...") train_input_names, train_output_names, val_input_names, val_output_names, test_input_names, test_output_names = utils.prepare_data( dataset_dir=args["dataset"]) _log.info("\n***** Begin training *****") _log.debug("Dataset -->", args["dataset"]) _log.debug("Model -->", args["model"]) _log.debug("Crop Height -->", args["crop_height"]) _log.debug("Crop Width -->", args["crop_width"]) _log.debug("Num Epochs -->", args["num_epochs"]) _log.debug("Batch Size -->", args["batch_size"]) _log.debug("Num Classes -->", num_classes) _log.debug("Data Augmentation:") _log.debug("\tVertical Flip -->", args["v_flip"]) _log.debug("\tHorizontal Flip -->", args["h_flip"]) _log.debug("\tBrightness Alteration -->", args["brightness"]) _log.debug("\tRotation -->", args["rotation"]) avg_loss_per_epoch = [] avg_scores_per_epoch = [] avg_iou_per_epoch = [] # Which validation images do we want val_indices = [] num_vals = min(args["num_val_images"], len(val_input_names)) # Set random seed to make sure models are validated on the same validation images. # So you can compare the results of different models more intuitively. random.seed(16) val_indices = random.sample(range(0, len(val_input_names)), num_vals) # Do the training here for epoch in range(args["epoch_start_i"], args["num_epochs"]): current_losses = [] cnt = 0 # Equivalent to shuffling id_list = np.random.permutation(len(train_input_names)) num_iters = int(np.floor(len(id_list) / args["batch_size"])) st = time.time() epoch_st = time.time() for i in range(num_iters): # st=time.time() input_image_batch = [] output_image_batch = [] # Collect a batch of images for j in range(args["batch_size"]): index = i * args["batch_size"] + j id = id_list[index] input_image = utils.load_image(train_input_names[id], args["crop_width"], args["crop_height"]) output_image = utils.load_image(train_output_names[id], args["crop_width"], args["crop_height"]) with tf.device('/cpu:0'): input_image, output_image = data_augmentation( input_image, output_image) # Prep the data. Make sure the labels are in one-hot format input_image = np.float32(input_image) / 255.0 output_image = np.float32( helpers.one_hot_it(label=output_image, label_values=label_values)) input_image_batch.append( np.expand_dims(input_image, axis=0)) output_image_batch.append( np.expand_dims(output_image, axis=0)) if args["batch_size"] == 1: input_image_batch = input_image_batch[0] output_image_batch = output_image_batch[0] else: input_image_batch = np.squeeze( np.stack(input_image_batch, axis=1)) output_image_batch = np.squeeze( np.stack(output_image_batch, axis=1)) # Do the training _, current = sess.run([opt, loss], feed_dict={ net_input: input_image_batch, net_output: output_image_batch }) current_losses.append(current) cnt = cnt + args["batch_size"] if cnt % 20 == 0: string_print = "Epoch = %d Count = %d Iter:(%d of %d) Current_Loss = %.4f Time = %.2f" % ( epoch, cnt, i, num_iters, current, time.time() - st) utils.LOG(string_print) st = time.time() mean_loss = np.mean(current_losses) avg_loss_per_epoch.append(mean_loss) # Create directories if needed if not os.path.isdir("%s/%04d" % ("checkpoints", epoch)): os.makedirs("%s/%04d" % ("checkpoints", epoch)) # Save latest checkpoint to same file name _log.info("Saving latest checkpoint") saver.save(sess, model_checkpoint_name) if val_indices != 0 and epoch % args["checkpoint_step"] == 0: _log.info("Saving checkpoint for this epoch") saver.save(sess, "%s/%04d/model.ckpt" % ("checkpoints", epoch)) if epoch % args["validation_step"] == 0: _log.info("Performing validation") target_path = "%s/%04d/val_scores.csv" % ("checkpoints", epoch) target = open(target_path, 'w') target.write( "val_name, avg_accuracy, precision, recall, f1 score, mean iou, %s\n" % (class_names_string)) scores_list = [] class_scores_list = [] precision_list = [] recall_list = [] f1_list = [] iou_list = [] # Do the validation on a small set of validation images for ind in val_indices: input_image = np.expand_dims(np.float32( utils.load_image(val_input_names[ind], args["crop_width"], args["crop_height"]) [:args["crop_height"], :args["crop_width"]]), axis=0) / 255.0 gt = utils.load_image( val_output_names[ind], args["crop_width"], args["crop_height"] )[:args["crop_height"], :args["crop_width"]] gt = helpers.reverse_one_hot( helpers.one_hot_it(gt, label_values)) # st = time.time() output_image = sess.run(network, feed_dict={net_input: input_image}) output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation( output_image, label_values) accuracy, class_accuracies, prec, rec, f1, iou = utils.evaluate_segmentation( pred=output_image, label=gt, num_classes=num_classes) file_name = utils.filepath_to_name(val_input_names[ind]) target.write("%s, %f, %f, %f, %f, %f" % (file_name, accuracy, prec, rec, f1, iou)) for item in class_accuracies: target.write(", %f" % (item)) target.write("\n") scores_list.append(accuracy) class_scores_list.append(class_accuracies) precision_list.append(prec) recall_list.append(rec) f1_list.append(f1) iou_list.append(iou) gt = helpers.colour_code_segmentation(gt, label_values) file_name = os.path.basename(val_input_names[ind]) file_name = os.path.splitext(file_name)[0] pred_img_path = "%s/%04d/%s_pred.png" % ("checkpoints", epoch, file_name) gt_img_path = "%s/%04d/%s_gt.png" % ("checkpoints", epoch, file_name) cv2.imwrite( pred_img_path, cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) cv2.imwrite(gt_img_path, cv2.cvtColor(np.uint8(gt), cv2.COLOR_RGB2BGR)) # Send the first 16 images to sacred if ind in val_indices[:16]: ex.add_artifact(gt_img_path, "GtImage_%d" % ind) ex.add_artifact(pred_img_path, "PredImage_%d" % ind) target.close() ex.add_artifact(target_path) avg_score = np.mean(scores_list) class_avg_scores = np.mean(class_scores_list, axis=0) avg_scores_per_epoch.append(avg_score) avg_precision = np.mean(precision_list) avg_recall = np.mean(recall_list) avg_f1 = np.mean(f1_list) avg_iou = np.mean(iou_list) avg_iou_per_epoch.append(avg_iou) # Sacred info dict gets sent every heartbeat (10s) ex.info["avg_score"] = avg_score ex.info["class_avg_scores"] = class_avg_scores ex.info["avg_precision"] = avg_precision ex.info["avg_recall"] = avg_recall ex.info["avg_f1"] = avg_f1 ex.info["avg_iou"] = avg_iou _log.debug("\nAverage validation accuracy for epoch # %04d = %f" % (epoch, avg_score)) _log.debug( "Average per class validation accuracies for epoch # %04d:" % (epoch)) for index, item in enumerate(class_avg_scores): _log.debug("%s = %f" % (class_names_list[index], item)) _log.debug("Validation precision = ", avg_precision) _log.debug("Validation recall = ", avg_recall) _log.debug("Validation F1 score = ", avg_f1) _log.debug("Validation IoU score = ", avg_iou) epoch_time = time.time() - epoch_st remain_time = epoch_time * (args["num_epochs"] - 1 - epoch) m, s = divmod(remain_time, 60) h, m = divmod(m, 60) if s != 0: train_time = "Remaining training time = %d hours %d minutes %d seconds\n" % ( h, m, s) else: train_time = "Remaining training time : Training completed.\n" utils.LOG(train_time) scores_list = [] fig1, ax1 = plt.subplots(figsize=(11, 8)) ax1.plot(range(epoch + 1), avg_scores_per_epoch) ax1.set_title("Average validation accuracy vs epochs") ax1.set_xlabel("Epoch") ax1.set_ylabel("Avg. val. accuracy") plt.savefig('accuracy_vs_epochs.png') ex.add_artifact("accuracy_vs_epochs.png") plt.clf() fig2, ax2 = plt.subplots(figsize=(11, 8)) ax2.plot(range(epoch + 1), avg_loss_per_epoch) ax2.set_title("Average loss vs epochs") ax2.set_xlabel("Epoch") ax2.set_ylabel("Current loss") plt.savefig('loss_vs_epochs.png') ex.add_artifact("loss_vs_epochs.png") plt.clf() fig3, ax3 = plt.subplots(figsize=(11, 8)) ax3.plot(range(epoch + 1), avg_iou_per_epoch) ax3.set_title("Average IoU vs epochs") ax3.set_xlabel("Epoch") ax3.set_ylabel("Current IoU") plt.savefig('iou_vs_epochs.png') ex.add_artifact("iou_vs_epochs.png")
sess.run(tf.global_variables_initializer()) print('Loading model checkpoint weights') saver=tf.train.Saver(max_to_keep=1000) saver.restore(sess, args.checkpoint_path) for ind in pre_indices: print("Testing image " + pre_input_names[ind]) loaded_image = utils.load_image(pre_input_names[ind]) padding_image = cv2.copyMakeBorder(loaded_image, 14, 13, 14, 13, cv2.BORDER_REFLECT) input_image = np.expand_dims(np.float32(padding_image),axis=0)/255.0 st = time.time() output_image = sess.run(network,feed_dict={net_input:input_image}) run_time = time.time()-st output_image = np.array(output_image[0,:,:,:]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation(output_image, label_values)[14:-13,14:-13,:] file_name = utils.filepath_to_name(pre_input_names[ind]) if not os.path.isdir("%s/predict/%s"%(args.dataset,args.model)): os.makedirs("%s/predict/%s"%(args.dataset,args.model)) cv2.imwrite("%s/predict/%s/%s_pred.png"%(args.dataset,args.model,file_name),cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) print("Wrote image " + "%s/predict/%s/%s_pred.png"%(args.dataset,args.model,file_name)) print("") print("Finished!")
def segment(file_list): parser = argparse.ArgumentParser() image = file_list[0] checkpoint_path = "latest_model_DeepLabV3_plus_CamVid_people.ckpt" crop_height = 512 crop_width = 512 model = "DeepLabV3_plus" dataset = "CamVid_people" class_names_list, label_values = helpers.get_label_info( os.path.join(args.dataset, "class_dict.csv")) num_classes = len(label_values) #num_classes = 2 print("\n***** Begin prediction *****") print("Dataset -->", args.dataset) print("Model -->", args.model) print("Crop Height -->", args.crop_height) print("Crop Width -->", args.crop_width) print("Num Classes -->", num_classes) print("Image -->", args.image) # Initializing network config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3]) net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes]) network, _ = model_builder.build_model(net_input=net_input, num_classes=num_classes, is_training=False) sess.run(tf.global_variables_initializer()) print('Loading model checkpoint weights') saver = tf.train.Saver(max_to_keep=1000) saver.restore(sess, args.checkpoint_path) print("Testing image " + args.image) #loaded_image = utils.load_image(args.image) images = [ "./CamVid_people/test/0001TP_007170.png", "./CamVid_people/test/0001TP_007860.png", "./CamVid_people/test/0001TP_008070.png" ] for image in images: loaded_image = utils.load_image(image) resized_image = cv2.resize(loaded_image, (args.crop_width, args.crop_height)) input_image = np.expand_dims(np.float32( resized_image[:args.crop_height, :args.crop_width]), axis=0) / 255.0 st = time.time() output_image = sess.run(network, feed_dict={net_input: input_image}) run_time = time.time() - st output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation( output_image, label_values) file_name = utils.filepath_to_name(args.image) cv2.imwrite("%s_pred.png" % (file_name), cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) print("") print("Finished!") print("time elapsed", run_time) print("Wrote image " + "%s_pred.png" % (file_name))
(args.crop_width, args.crop_height)) input_image = np.expand_dims(np.float32( resized_image[:args.crop_height, :args.crop_width]), axis=0) / 255.0 st = time.time() output_image = sess.run(network, feed_dict={net_input: input_image}) run_time = time.time() - st output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation(output_image, label_values) overlayed_img = 0.35 * resized_image + 0.65 * out_vis_image # overlayed_img = overlayed_img.astype(np.uint8) file_name = utils.filepath_to_name(image) cv2.imwrite("preds/%s_pred.png" % (file_name), cv2.cvtColor(np.uint8(overlayed_img), cv2.COLOR_RGB2BGR)) cv2.imwrite( "preds/%s_mask.png" % (file_name), cv2.cvtColor((cv2.resize(np.uint8(out_vis_image), (img_shape[0], img_shape[1]))), cv2.COLOR_RGB2BGR)) print("") print("Finished!") print("Wrote image " + "%s_pred.png" % (file_name))
def main(args=None): #If args is None, will parse command line arguments #Otherwise args needs to be list with arguments, like # ['--dataset', 'CamVid', '--model', 'FC-DenseNet103'] parser = argparse.ArgumentParser() parser.add_argument('--num_epochs', type=int, default=300, help='Number of epochs to train for') parser.add_argument('--epoch_start_i', type=int, default=0, help='Start counting epochs from this number') parser.add_argument('--checkpoint_step', type=int, default=5, help='How often to save checkpoints (epochs)') parser.add_argument('--validation_step', type=int, default=1, help='How often to perform validation (epochs)') parser.add_argument('--continue_training', type=utils.str2bool, default=False, help='Whether to continue training from a checkpoint') parser.add_argument( '--continue_from', type=str, default='', help='From which checkpoint to continue. Only relevant with darea_call.' ) parser.add_argument('--dataset', type=str, default="CamVid", help='Dataset you are using.') parser.add_argument('--dataset_path', type=str, default="", help='Path to Dataset folder.') parser.add_argument( '--image_suffix', type=str, default='', required=False, help= 'Only files with this extension should be included. You should specify it if some non-image files will be in the same folder' ) parser.add_argument('--crop_height', type=int, default=512, help='Height of cropped input image to network') parser.add_argument('--crop_width', type=int, default=512, help='Width of cropped input image to network') parser.add_argument( '--biased_crop', type=float, default=0, help= 'Probability of making a biased cropped. Biased crops always contain some foreground portion. Only works if one of the classes is named "Background".' ) parser.add_argument( '--downscale_factor', type=float, default=0, required=False, help= 'Shrink image by this factor. E.g. if image is 1024x1024 and downscale_factor is 0.5, downscaled image will be 512x512. This is applied before cropping.' ) parser.add_argument('--batch_size', type=int, default=1, help='Number of images in each batch') parser.add_argument( '--num_val_images', type=int, default=20, help='The number of images to used for validations. If -1 -> use all') parser.add_argument( '--h_flip', type=utils.str2bool, default=False, help= 'Whether to randomly flip the image horizontally for data augmentation' ) parser.add_argument( '--v_flip', type=utils.str2bool, default=False, help= 'Whether to randomly flip the image vertically for data augmentation') parser.add_argument( '--brightness', type=float, default=None, help= 'Whether to randomly change the image brightness for data augmentation. Specifies the max bightness change as a factor between 0.0 and 1.0. For example, 0.1 represents a max brightness change of 10%% (+-).' ) parser.add_argument( '--rotation', type=float, default=None, help= 'DOES NOT WORK! Whether to randomly rotate the image for data augmentation. Specifies the max rotation angle in degrees.' ) parser.add_argument('--rotation_perpendicular', type=utils.str2bool, default=False, help='Randomly rotates by 0, 90, 180 or 270 degrees') parser.add_argument( '--model', type=str, default="FC-DenseNet103", help= 'The model you are using. See model_builder.py for supported models') parser.add_argument( '--frontend', type=str, default="ResNet101", help= 'The frontend you are using. See frontend_builder.py for supported models' ) parser.add_argument( '--save_best', type=utils.str2bool, default=False, help='Saves model with smallest loss rather than last model') parser.add_argument('--learn_rate', type=float, default=0.0001, help='The learning rate') parser.add_argument( '--chkpt_prefix', type=str, default='', help= 'Prefix in front of checkpoint (intended for distinguishing same models ran with different parameters)' ) parser.add_argument( '--darea_call', type=utils.str2bool, default=False, required=False, help='Set to true when you call it from Darea software') parser.add_argument( '--save_path', type=str, default='checkpoints', required=False, help='Name of saved model. Only used when darea_call is true.') parser.add_argument('--makePlots', type=utils.str2bool, default=True, required=False, help='Whether plots should be made') if args is None: args = parser.parse_args() else: args = parser.parse_args(args) if args.darea_call: os.chdir(os.path.dirname(os.path.realpath(__file__))) # Get the names of the classes so we can record the evaluation results class_names_list, label_values = helpers.get_label_info( os.path.join(args.dataset_path, args.dataset, "class_dict.csv")) class_names_string = ', '.join(class_names_list) if 'Background' not in class_names_list: args.biased_crop = 0 backgroundValue = None else: backgroundValue = label_values[class_names_list.index('Background')] num_classes = len(label_values) config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) # Compute your softmax cross entropy loss net_input = tf.placeholder(tf.float32, shape=[ None, None, None, 3 ]) #Try setting to 1 for grayscale, think theres no other changes needed. net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes]) network, init_fn = model_builder.build_model(model_name=args.model, frontend=args.frontend, net_input=net_input, num_classes=num_classes, crop_width=args.crop_width, crop_height=args.crop_height, is_training=True) loss = tf.reduce_mean( tf.nn.softmax_cross_entropy_with_logits_v2(logits=network, labels=net_output)) opt = tf.train.RMSPropOptimizer( learning_rate=args.learn_rate, decay=0.995).minimize( loss, var_list=[var for var in tf.trainable_variables()]) saver = tf.train.Saver(max_to_keep=500) sess.run(tf.global_variables_initializer()) utils.count_params() # If a pre-trained ResNet is required, load the weights. # This must be done AFTER the variables are initialized with sess.run(tf.global_variables_initializer()) if init_fn is not None: init_fn(sess) # Load a previous checkpoint if desired if args.darea_call: if args.continue_training and args.continue_from: model_checkpoint_name = '../../deepLearning/checkpoints/' + args.continue_from + '.ckpt' if os.path.isfile(model_checkpoint_name + '.index'): print('Loading latest model checkpoint...') saver.restore(sess, model_checkpoint_name) print('successfully loaded {}'.format(model_checkpoint_name)) else: print('Specified checkpoint {} not found. Starting fresh one'. format(os.path.abspath(model_checkpoint_name))) if args.save_path: model_checkpoint_name = os.path.join(args.save_path, args.dataset + '.ckpt') else: if args.continue_training: if args.continue_from: model_checkpoint_name = args.continue_from if not model_checkpoint_name.endswith('.ckpt'): model_checkpoint_name += '.ckpt' else: model_checkpoint_name = os.path.join( args.save_path, "latest_model_" + args.chkpt_prefix + args.model + "_" + args.dataset + ".ckpt") if os.path.isfile(model_checkpoint_name + '.index'): print('Loading latest model checkpoint...') saver.restore(sess, model_checkpoint_name) print('successfully loaded {}'.format(model_checkpoint_name)) else: print('{} not found. Starting a fresh training'.format( args.continue_from)) model_checkpoint_name = os.path.join( args.save_path, "latest_model_" + args.chkpt_prefix + args.model + "_" + args.dataset + ".ckpt") # Load the data print("Loading the data ...") train_input_names,train_output_names, val_input_names, val_output_names, test_input_names, test_output_names = \ utils.prepare_data(dataset_dir=os.path.join(args.dataset_path,args.dataset),image_suffix=args.image_suffix) print("\n***** Begin training *****") print("Dataset -->", args.dataset) print("Model -->", args.model) if args.downscale_factor: print("Downscale Factor -->", args.downscale_factor) print("Crop Height -->", args.crop_height) print("Crop Width -->", args.crop_width) print("Num Epochs -->", args.num_epochs) print("Batch Size -->", args.batch_size) print("Num Classes -->", num_classes) print("Learn Rate -->", args.learn_rate) print("Validation Step -->", args.validation_step) print("Data Augmentation:") print("\tVertical Flip -->", args.v_flip) print("\tHorizontal Flip -->", args.h_flip) print("\tBrightness Alteration -->", args.brightness) print("\tRotation -->", args.rotation) print("\tPerpendicular Rotation -->", args.rotation_perpendicular) print("", flush=True) avg_loss_per_epoch = [] avg_scores_per_epoch = [] avg_iou_per_epoch = [] if args.num_val_images == -1: #Use all validation images if so wanted by users val_indices = range(0, len(val_input_names)) else: # Which validation images do we want val_indices = [] num_vals = min(args.num_val_images, len(val_input_names)) # Set random seed to make sure models are validated on the same validation images. # So you can compare the results of different models more intuitively. random.seed(16) val_indices = random.sample(range(0, len(val_input_names)), num_vals) #Copy class file with same name as checkpoint shutil.copyfile( os.path.join(args.dataset_path, args.dataset, "class_dict.csv"), os.path.splitext(model_checkpoint_name)[0] + '.classes') # Do the training here for epoch in range(args.epoch_start_i, args.num_epochs): current_losses = [] cnt = 0 # Equivalent to shuffling id_list = np.random.permutation(len(train_input_names)) num_iters = int(np.floor(len(id_list) / args.batch_size)) st = time.time() epoch_st = time.time() for i in range(num_iters): # st=time.time() input_image_batch = [] output_image_batch = [] # Collect a batch of images for j in range(args.batch_size): index = i * args.batch_size + j id = id_list[index] input_image = utils.load_image(train_input_names[id]) output_image = utils.load_image(train_output_names[id]) with tf.device('/cpu:0'): input_image, output_image = data_augmentation( input_image, output_image, args, backgroundValue) if 0: #Debugging: try: os.mkdir('imagesinTrain') 'kj' except: pass try: os.mkdir('imagesinTrain/%04d' % (epoch)) except: pass file_name = utils.filepath_to_name( train_input_names[id]) file_name = os.path.splitext(file_name)[0] cv2.imwrite( "%s/%04d/%s_im.png" % ("imagesinTrain", epoch, file_name), cv2.cvtColor(np.uint8(input_image), cv2.COLOR_RGB2BGR)) cv2.imwrite( "%s/%04d/%s_gt.png" % ("imagesinTrain", epoch, file_name), cv2.cvtColor(np.uint8(output_image), cv2.COLOR_RGB2BGR)) # Prep the data. Make sure the labels are in one-hot format input_image = np.float32(input_image) / 255.0 output_image = np.float32( helpers.one_hot_it(label=output_image, label_values=label_values)) if 0: #Debugging: try: os.mkdir('imagesinTrain') except: pass try: os.mkdir('imagesinTrain/%04d' % (epoch)) except: pass file_name = utils.filepath_to_name( train_input_names[id]) file_name = os.path.splitext(file_name)[0] print("%s/%04d/%s_im.png" % ("imagesinTrain", epoch, file_name)) cv2.imwrite( "%s/%04d/%s_im.png" % ("imagesinTrain", epoch, file_name), cv2.cvtColor(np.uint8(input_image), cv2.COLOR_RGB2BGR)) #cv2.imwrite("%s/%04d/%s_gt.png"%("imagesinTrain",epoch, file_name),cv2.cvtColor(np.uint8(output_image), cv2.COLOR_RGB2BGR)) input_image_batch.append( np.expand_dims(input_image, axis=0)) output_image_batch.append( np.expand_dims(output_image, axis=0)) if args.batch_size == 1: input_image_batch = input_image_batch[0] output_image_batch = output_image_batch[0] else: input_image_batch = np.squeeze( np.stack(input_image_batch, axis=1)) output_image_batch = np.squeeze( np.stack(output_image_batch, axis=1)) # Do the training _, current, output_image = sess.run([opt, loss, network], feed_dict={ net_input: input_image_batch, net_output: output_image_batch }) current_losses.append(current) cnt = cnt + args.batch_size if cnt % 25 == 0: string_print = "Epoch = %d Count = %d Current_Loss = %.4f Time = %.2f" % ( epoch, cnt, current, time.time() - st) utils.LOG(string_print) st = time.time() if 0: #For Debugging output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation( output_image, label_values) cv2.imwrite( "%s/%04d/%s_pred.png" % ("imagesinTrain", epoch, file_name), cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) mean_loss = np.mean(current_losses) avg_loss_per_epoch.append(mean_loss) # Create directories if needed if not os.path.isdir("%s/%04d" % ("checkpoints", epoch)): os.makedirs("%s/%04d" % ("checkpoints", epoch)) # If latest checkpoint should be saved, save now to same file name, if not args.save_best: print("Saving latest checkpoint") saver.save(sess, model_checkpoint_name) if val_indices != 0 and epoch % args.checkpoint_step == 0: print("Saving checkpoint for this epoch") saver.save(sess, "%s/%04d/model.ckpt" % (args.save_path, epoch)) if epoch % args.validation_step == 0: if epoch == 0: best_avg_iou = 0 print("Performing validation") target = open("%s/%04d/val_scores.csv" % ("checkpoints", epoch), 'w') target.write( "val_name, avg_accuracy, precision, recall, f1 score, mean iou, %s\n" % (class_names_string)) scores_list = [] class_scores_list = [] precision_list = [] recall_list = [] f1_list = [] iou_list = [] # Do the validation on a small set of validation images for ind in val_indices: input_image = utils.load_image(val_input_names[ind]) gt = utils.load_image(val_output_names[ind]) if args.downscale_factor and args.downscale_factor != 1: dim = (int(input_image.shape[0] * args.downscale_factor), int(input_image.shape[1] * args.downscale_factor)) input_image = cv2.resize(input_image, dim, interpolation=cv2.INTER_CUBIC) gt = cv2.resize(gt, dim, interpolation=cv2.INTER_NEAREST) #input_image, gt = data_augmentation(input_image, gt, args) gt = helpers.reverse_one_hot( helpers.one_hot_it(gt, label_values)) crop_height = args.crop_height crop_width = args.crop_width if input_image.shape[0] > crop_height or input_image.shape[ 1] > crop_width: #rectangle in bottom right corner smaller than cropped im will not be used nrCroppingsY = input_image.shape[0] // crop_height nrCroppingsX = input_image.shape[1] // crop_width output_image = np.zeros([ nrCroppingsY * crop_height, nrCroppingsX * crop_width ]) gt = gt[0:nrCroppingsY * crop_height, 0:nrCroppingsX * crop_width] for yi in range(nrCroppingsY): row = np.zeros( [crop_height, nrCroppingsX * crop_width]) for xi in range(nrCroppingsX): inputIm = input_image[yi * crop_height:(1 + yi) * crop_height, xi * crop_width:(1 + xi) * crop_width, :] inputIm = np.expand_dims(np.float32(inputIm) / 255.0, axis=0) out = sess.run(network, feed_dict={net_input: inputIm}) out = np.array(out[0, :, :, :]) out = helpers.reverse_one_hot(out) row[:, xi * crop_width:(1 + xi) * crop_width] = out output_image[yi * crop_height:(1 + yi) * crop_height, :] = row # st = time.time() else: input_image = np.expand_dims(np.float32(input_image) / 255.0, axis=0) output_image = sess.run(network, feed_dict={net_input: input_image}) output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation( output_image, label_values) accuracy, class_accuracies, prec, rec, f1, iou, class_iou = utils.evaluate_segmentation( pred=output_image, label=gt, num_classes=num_classes) file_name = utils.filepath_to_name(val_input_names[ind]) target.write("%s, %f, %f, %f, %f, %f" % (file_name, accuracy, prec, rec, f1, iou)) for item in class_accuracies: target.write(", %f" % (item)) target.write("\n") scores_list.append(accuracy) class_scores_list.append(class_accuracies) precision_list.append(prec) recall_list.append(rec) f1_list.append(f1) iou_list.append(iou) gt = helpers.colour_code_segmentation(gt, label_values) file_name = os.path.basename(val_input_names[ind]) file_name = os.path.splitext(file_name)[0] cv2.imwrite( "%s/%04d/%s_pred.png" % ("checkpoints", epoch, file_name), cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) cv2.imwrite( "%s/%04d/%s_gt.png" % ("checkpoints", epoch, file_name), cv2.cvtColor(np.uint8(gt), cv2.COLOR_RGB2BGR)) target.close() avg_score = np.mean(scores_list) class_avg_scores = np.mean(class_scores_list, axis=0) avg_scores_per_epoch.append(avg_score) avg_precision = np.mean(precision_list) avg_recall = np.mean(recall_list) avg_f1 = np.mean(f1_list) avg_iou = np.mean(iou_list) avg_iou_per_epoch.append(avg_iou) print("\nAverage validation accuracy for epoch # %04d = %f" % (epoch, avg_score)) print("Average per class validation accuracies for epoch # %04d:" % (epoch)) for index, item in enumerate(class_avg_scores): print("%s = %f" % (class_names_list[index], item)) print("Validation precision = ", avg_precision) print("Validation recall = ", avg_recall) print("Validation F1 score = ", avg_f1) print("Validation IoU score = ", avg_iou, flush=True) if args.save_best and avg_iou > best_avg_iou: #Save model if best model is wanted and it is the best or if it is first evaluation print("Saving best checkpoint with iou {}".format(avg_iou)) saver.save(sess, model_checkpoint_name) best_avg_iou = avg_iou #Save an info file with open(model_checkpoint_name[:-5] + '.info', 'w') as f: f.write('Epoch\t{}\nValidation IoU score\t{}'.format( epoch, avg_iou)) epoch_time = time.time() - epoch_st remain_time = epoch_time * (args.num_epochs - 1 - epoch) m, s = divmod(remain_time, 60) h, m = divmod(m, 60) if s != 0: train_time = "Remaining training time = %d hours %d minutes %d seconds\n" % ( h, m, s) else: train_time = "Remaining training time : Training completed.\n" utils.LOG(train_time) scores_list = [] with open(model_checkpoint_name[:-5] + '.info', 'a+') as f: #Save some info on filesizes f.write('\nimageSize\t{}'.format([args.crop_height, args.crop_width])) f.write('\nTraining completed\t{}'.format( datetime.datetime.now().strftime("%d-%b-%Y %H:%M:%S"))) if not args.darea_call and args.makePlots: fig1, ax1 = plt.subplots(figsize=(11, 8)) ax1.plot(range(epoch + 1), avg_scores_per_epoch) ax1.set_title("Average validation accuracy vs epochs") ax1.set_xlabel("Epoch") ax1.set_ylabel("Avg. val. accuracy") plt.savefig('accuracy_vs_epochs.png') plt.clf() fig2, ax2 = plt.subplots(figsize=(11, 8)) ax2.plot(range(epoch + 1), avg_loss_per_epoch) ax2.set_title("Average loss vs epochs") ax2.set_xlabel("Epoch") ax2.set_ylabel("Current loss") plt.savefig('loss_vs_epochs.png') plt.clf() fig3, ax3 = plt.subplots(figsize=(11, 8)) ax3.plot(range(epoch + 1), avg_iou_per_epoch) ax3.set_title("Average IoU vs epochs") ax3.set_xlabel("Epoch") ax3.set_ylabel("Current IoU") plt.savefig('iou_vs_epochs.png')
def main(args=None): #If args is None, will parse command line arguments #Otherwise args needs to be list with arguments, like # ['--image', 'imageFolder', '--model', 'FC-DenseNet103'] parser = argparse.ArgumentParser() parser.add_argument( '--image', type=str, default=None, required=True, help='The image or folder with images you want to predict on. ') parser.add_argument( '--checkpoint_path', type=str, default=None, required=False, help= 'The path to the latest checkpoint weights for your model. Will guess based on model and dataset if not specified' ) parser.add_argument('--crop_height', type=int, default=512, help='Height of cropped input image to network') parser.add_argument('--crop_width', type=int, default=512, help='Width of cropped input image to network') parser.add_argument( '--downscale_factor', type=float, default=0, required=False, help= 'Shrink image by this factor. E.g. if image is 1024x1024 and downscale_factor is 0.5, downscaled image will be 512x512. This is applied before cropping.' ) parser.add_argument('--model', type=str, default="FC-DenseNet103", required=False, help='The model you are using') parser.add_argument('--dataset', type=str, default="CamVid", required=False, help='The dataset you are using') parser.add_argument( '--image_suffix', type=str, default='', required=False, help= 'Only files with this extension should be included. You should specify it if some non-image files will be in the same folder' ) parser.add_argument( '--outpath', type=str, default='./', required=False, help='Folder where predicted images should be saved to') parser.add_argument( '--file_suffix', type=str, default='_pred.png', required=False, help='Suffix appended to input image name for output image') parser.add_argument( '--darea_call', type=utils.str2bool, default=False, required=False, help='Set to true when you call it from Darea software') parser.add_argument('--save_predictionImage', type=utils.str2bool, default=True, required=False, help='Whether predictions should be saved as images') parser.add_argument( '--save_coordinates', type=utils.str2bool, default=False, required=False, help='Whether coordinates of predicted structures should be saved') if args is None: args = parser.parse_args() else: args = parser.parse_args(args) if args.darea_call: os.chdir(os.path.dirname(os.path.realpath(__file__))) class_names_list, label_values = helpers.get_label_info( os.path.join(args.dataset, "class_dict.csv")) num_classes = len(label_values) print("\n***** Begin prediction *****") print("Dataset -->", args.dataset) print("Model -->", args.model) print("Crop Height -->", args.crop_height) print("Crop Width -->", args.crop_width) print("Num Classes -->", num_classes) print("Image -->", args.image) if args.checkpoint_path is None: args.checkpoint_path = "checkpoints/latest_model_" + args.model + "_" + args.dataset + ".ckpt" # Initializing network config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3]) net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes]) network, _ = model_builder.build_model(args.model, net_input=net_input, num_classes=num_classes, crop_width=args.crop_width, crop_height=args.crop_height, is_training=False) sess.run(tf.global_variables_initializer()) print('Loading model checkpoint weights') saver = tf.train.Saver(max_to_keep=1000) saver.restore(sess, args.checkpoint_path) if os.path.isdir(args.image): folder = args.image #Get list of all images in the folder, ignore hidden files, subfolders and files not ending with correct suffix images = [ os.path.join(folder, file) for file in os.listdir(folder) if not file.startswith('.') and not os.path.isdir(file) and file.endswith(args.image_suffix) ] else: images = [args.image] if not os.path.isdir(args.outpath): os.mkdir(args.outpath) if args.save_coordinates: coordinates = dict() print('Performing predictions...\n') progress_string = '' for index, image in enumerate(images): if args.darea_call: if index > 1: print('\b' * (len(progress_string) + 2)) #Deletes previous output in matlab console progress_string = 'Predicting on image {} / {}'.format( index + 1, len(images)) print(progress_string) else: print('Predicting on image {} / {}\r'.format( index + 1, len(images))) input_image = utils.load_image(image) if args.downscale_factor and args.downscale_factor != 1: dim = (int(input_image.shape[0] * args.downscale_factor), int(input_image.shape[1] * args.downscale_factor)) input_image = cv2.resize(input_image, dim, interpolation=cv2.INTER_CUBIC) st = time.time() crop_height = args.crop_height crop_width = args.crop_width if input_image.shape[0] > crop_height or input_image.shape[ 1] > crop_width: #rectangle in bottom right corner smaller than cropped im will not be used nrCroppingsY = m.ceil(input_image.shape[0] / crop_height) nrCroppingsX = m.ceil(input_image.shape[1] / crop_width) output_image = np.zeros( [nrCroppingsY * crop_height, nrCroppingsX * crop_width]) for yi in range(nrCroppingsY): row = np.zeros([crop_height, nrCroppingsX * crop_width]) cropYstart = yi * crop_height cropYstop = (1 + yi) * crop_height if cropYstop >= input_image.shape[0]: cropYstop = input_image.shape[0] - 1 cropYstart = cropYstop - crop_height for xi in range(nrCroppingsX): cropXstart = xi * crop_width cropXstop = (1 + xi) * crop_width if cropXstop >= input_image.shape[1]: cropXstop = input_image.shape[1] - 1 cropXstart = cropXstop - crop_width inputIm = input_image[cropYstart:cropYstop, cropXstart:cropXstop, :] inputIm = np.expand_dims(np.float32(inputIm) / 255.0, axis=0) #print(inputIm.shape) out = sess.run(network, feed_dict={net_input: inputIm}) out = np.array(out[0, :, :, :]) out = helpers.reverse_one_hot(out) if (1 + xi) * crop_width >= input_image.shape[1]: row[:, xi * crop_width:] = out else: row[:, xi * crop_width:(1 + xi) * crop_width] = out if (1 + yi) * crop_height >= input_image.shape[0]: output_image[yi * crop_height:, :] = row else: output_image[yi * crop_height:(1 + yi) * crop_height, :] = row # st = time.time() else: input_image = np.expand_dims(np.float32(input_image) / 255.0, axis=0) output_image = sess.run(network, feed_dict={net_input: input_image}) output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) run_time = time.time() - st out_vis_image = helpers.colour_code_segmentation( output_image, label_values) if args.save_predictionImage: file_name = "{}{}".format(utils.filepath_to_name(image), args.file_suffix) cv2.imwrite( os.path.join(args.outpath, file_name), cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) if args.save_coordinates: coordinates[utils.filepath_to_name( image)] = utils.getCoordsFromPrediction( cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2GRAY), args.downscale_factor) if not args.darea_call: print("Wrote image {}".format(file_name)) print("") if args.save_coordinates: with open(os.path.join(args.outpath, 'centroids.json'), 'w') as f: json.dump(coordinates, f) if not args.darea_call: print("Finished!")
def main(args=None): #If args is None, will parse command line arguments #Otherwise args needs to be list with arguments, like # ['--image', 'imageFolder', '--model', 'FC-DenseNet103'] parser = argparse.ArgumentParser() parser.add_argument('--image', type=str, default=None, required=True, help='The image or folder with images you want to predict on. ') parser.add_argument('--classes', type=str, required=False, default=None, help='Path to class_dict.csv. If None, expects it to be same as .ckpt but with .classes extension') parser.add_argument('--checkpoint_path', type=str, required=True, help='The path to the latest checkpoint weights for your model.') parser.add_argument('--crop_height', type=int, default=512, help='Height of cropped input image to network') parser.add_argument('--crop_width', type=int, default=512, help='Width of cropped input image to network') parser.add_argument('--downscale_factor', type=float, default=0, required=False, help='Shrink image by this factor. E.g. if image is 1024x1024 and downscale_factor is 0.5, downscaled image will be 512x512. This is applied before cropping.') parser.add_argument('--model', type=str, default="FC-DenseNet103", required=False, help='The model you are using') parser.add_argument('--image_suffix', type=str, default='', required=False, help='Only files with this extension should be included. You should specify it if some non-image files will be in the same folder') parser.add_argument('--outpath', type=str, default='./', required=False, help='Folder where predicted images should be saved to') parser.add_argument('--file_suffix', type=str, default='_pred.png', required=False, help='Suffix appended to input image name for output image') parser.add_argument('--save_predictionImage', type=utils.str2bool, default=True, required=False, help='Whether predictions should be saved as images') parser.add_argument('--save_coordinates', type=utils.str2bool, default=True, required=False, help='Whether coordinates of predicted structures should be saved') parser.add_argument('--coords_filename', type=str, default='coordinates',required=False, help='Filename of the coordinates file (without filextension).') parser.add_argument('--coords_class', type=str, default=None, required=False, help='Specify class to generate coordinates from. (Default is all classes). Separate multiple classes by comma to pool them.') parser.add_argument('--wait_time', type=float, default=1, required=False, help='Time (in s) to wait for new images to appear in folder') if args is None: args=parser.parse_args() else: args=parser.parse_args(args) if args.classes is None: args.classes=os.path.splitext(args.checkpoint_path)[0]+'.classes' print(args.classes) class_names_list, label_values = helpers.get_label_info(args.classes) num_classes = len(label_values) assert os.path.isdir(args.image), "Argument image needs to be a folder, got {}".format(args.image) if args.save_coordinates: if args.coords_class is None: foregroundcolors=[c[0] for c in label_values[1:]] #Labels are in RGB, but we only need grayscale value else: target_classes=args.coords_class.split(',') foregroundcolors=[label_values[i][0] for i,c in enumerate(class_names_list) if c in target_classes] if len(foregroundcolors)<1: raise ValueError('None of the specified --coords_class found in class dict.\nSpecified classes were {}\nExisting classes are {}'.format(target_classes,class_names_list)) print("\n***** Begin prediction *****") #print("Dataset -->", args.dataset) print("Model -->", args.model) print("Crop Height -->", args.crop_height) print("Crop Width -->", args.crop_width) print("Num Classes -->", num_classes) print("Image -->", args.image) if args.checkpoint_path is None: args.checkpoint_path = "checkpoints/latest_model_" + args.model + "_" + args.dataset + ".ckpt" # Initializing network config = tf.ConfigProto() config.gpu_options.allow_growth = True sess=tf.Session(config=config) net_input = tf.placeholder(tf.float32,shape=[None,None,None,3]) net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes]) network, _ = model_builder.build_model(args.model, net_input=net_input, num_classes=num_classes, crop_width=args.crop_width, crop_height=args.crop_height, is_training=False) sess.run(tf.global_variables_initializer()) print('Loading model checkpoint weights') saver=tf.train.Saver(max_to_keep=1000) saver.restore(sess, args.checkpoint_path) #time.sleep(args.wait_time*3) folder=args.image finished_images=[] images=getUnfinishedImages(folder,finished_images,args.image_suffix) if not os.path.isdir(args.outpath): os.mkdir(args.outpath) if args.save_coordinates: coordinates=dict() print('Performing predictions...') while images: time.sleep(args.wait_time) for index,image in enumerate(images): print("Predicting image {} (image {} out of {})".format(image,index+1, len(images))) input_image = utils.load_image(image) if args.downscale_factor and args.downscale_factor !=1: dim=(int(input_image.shape[0]*args.downscale_factor), int(input_image.shape[1]*args.downscale_factor)) input_image=cv2.resize(input_image,dim,interpolation=cv2.INTER_CUBIC) st = time.time() crop_height=args.crop_height crop_width=args.crop_width if input_image.shape[0]>crop_height or input_image.shape[1]>crop_width: #rectangle in bottom right corner smaller than cropped im will not be used nrCroppingsY=m.ceil(input_image.shape[0]/crop_height) nrCroppingsX=m.ceil(input_image.shape[1]/crop_width) output_image=np.zeros([nrCroppingsY*crop_height,nrCroppingsX*crop_width]) for yi in range(nrCroppingsY): row=np.zeros([crop_height,nrCroppingsX*crop_width]) cropYstart=yi*crop_height; cropYstop=(1+yi)*crop_height if cropYstop>=input_image.shape[0]: cropYstop=input_image.shape[0]-1 cropYstart=cropYstop-crop_height for xi in range(nrCroppingsX): cropXstart=xi*crop_width; cropXstop=(1+xi)*crop_width if cropXstop>=input_image.shape[1]: cropXstop=input_image.shape[1]-1 cropXstart=cropXstop-crop_width inputIm=input_image[cropYstart:cropYstop,cropXstart:cropXstop,:] inputIm=np.expand_dims(np.float32(inputIm)/255.0,axis=0) #print(inputIm.shape) out=sess.run(network,feed_dict={net_input:inputIm}) out = np.array(out[0,:,:,:]) out=helpers.reverse_one_hot(out) if (1+xi)*crop_width>=input_image.shape[1]: row[:,xi*crop_width:]=out else: row[:,xi*crop_width:(1+xi)*crop_width]=out if (1+yi)*crop_height>=input_image.shape[0]: output_image[yi*crop_height:,:]=row else: output_image[yi*crop_height:(1+yi)*crop_height,:]=row # st = time.time() else: input_image=np.expand_dims(np.float32(input_image)/255.0,axis=0) output_image = sess.run(network,feed_dict={net_input:input_image}) output_image = np.array(output_image[0,:,:,:]) output_image = helpers.reverse_one_hot(output_image) run_time = time.time()-st out_vis_image = helpers.colour_code_segmentation(output_image, label_values) if args.save_predictionImage: file_name = "{}{}".format(utils.filepath_to_name(image),args.file_suffix) cv2.imwrite(os.path.join(args.outpath,file_name),cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) print("Wrote image {}".format(file_name)) if args.save_coordinates: coordinates[utils.filepath_to_name(image)]=utils.getCoordsFromPrediction(cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2GRAY),foregroundcolors,args.downscale_factor); print("", flush=True) #Finished predicting all images in list finished_images+=images #Now check if new images appeared in folder, if so, repeat, else wait, retry then finish images=getUnfinishedImages(folder,finished_images,args.image_suffix) if images: continue print("Waiting for new images", flush=True) time.sleep(args.wait_time) images=getUnfinishedImages(folder,finished_images,args.image_suffix) if not images: break if args.save_coordinates: with open(os.path.join(args.outpath,'{}.json'.format(args.coords_filename)), 'w') as f: json.dump(coordinates,f) print("Finished!")
run_times_list.append(time.time()-st) #print('Predicted image, took {}'.format(time.time()-st)) #st = time.time() output_image = np.array(output_image[0,:,:,:]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation(output_image, label_values) #print('Computed output images, took {}'.format(time.time()-st)) #st = time.time() accuracy, class_accuracies, prec, rec, f1, iou, class_iou = utils.evaluate_segmentation(pred=output_image, label=gt, num_classes=num_classes) #print('Computed metrics, took {}'.format(time.time()-st)) #st = time.time() file_name = utils.filepath_to_name(input_name) target.write("\n%s,%g,%g,%g,%g,%g"%(file_name, accuracy, prec, rec, f1, iou)) for item in class_accuracies: target.write(",%g"%(item)) for item in class_iou: target.write(",%g"%(item)) scores_list.append(accuracy) class_scores_list.append(class_accuracies) precision_list.append(prec) recall_list.append(rec) f1_list.append(f1) iou_list.append(iou) #print('Wrote metrics, took {}'.format(time.time()-st)) #st = time.time()
def predict_segment(images, save=False): checkpoint_path = "./weights/weight3/DeepLabV3_plus_CamVid_people_weight3.ckpt" model = "DeepLabV3_plus" dataset = "CamVid_people" crop_height = 512 crop_width = 512 class_names_list, label_values = helpers.get_label_info(os.path.join(dataset, "class_dict.csv")) num_classes = len(label_values) #num_classes = 2 print("\n***** Begin prediction *****") print("Dataset -->", dataset) print("Model -->", model) print("Crop Height -->", crop_height) print("Crop Width -->", crop_width) print("Num Classes -->", num_classes) print("Image -->", images) # Initializing network config = tf.ConfigProto() config.gpu_options.allow_growth = True sess=tf.Session(config=config) net_input = tf.placeholder(tf.float32,shape=[None,None,None,3]) net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes]) network, _ = model_builder.build_model(model, net_input=net_input, num_classes=num_classes, crop_width=crop_width, crop_height=crop_height, is_training=False) sess.run(tf.global_variables_initializer()) print('Loading model checkpoint weights') saver=tf.train.Saver(max_to_keep=1000) saver.restore(sess, checkpoint_path) print("Testing image", images) input_image = [cv2.resize(utils.load_image(image), (crop_width, crop_height))[:crop_height, :crop_width] for image in images] input_image = np.array(input_image, dtype = np.float32)/255.0 st = time.time() output_images = sess.run(network,feed_dict={net_input:input_image}) run_time = time.time()-st output_images = helpers.reverse_one_hot(output_images) out_vis_images = helpers.colour_binary_segmentation(output_images) if save: for i in range(len(images)): out_vis_image = np.array(out_vis_images[i,:,:]) #output_image = helpers.reverse_one_hot(output_image) #out_vis_image = helpers.colour_code_segmentation(output_image, label_values) image = images[i] file_name = utils.filepath_to_name(image) cv2.imwrite("%s_pred.png"%(file_name),cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)) print("Wrote image " + "%s_pred.png"%(file_name)) #output_images = helpers.reverse_one_hot(output_images) #out_vis_images = helpers.colour_binary_segmentation(output_images) print("") print("segment Finished!") print("time elapsed",run_time) return out_vis_images
is_training=False) sess.run(tf.global_variables_initializer()) print('Loading model checkpoint weights') saver=tf.train.Saver(max_to_keep=1000) saver.restore(sess, args.checkpoint_path) for image in os.listdir(args.input_folder): if not os.path.splitext(image)[-1] == '.png': continue image_path = os.path.join(args.input_folder, image) loaded_image = utils.load_image(image_path) resized_image =cv2.resize(loaded_image, (args.crop_width, args.crop_height)) input_image = np.expand_dims(np.float32(resized_image[:args.crop_height, :args.crop_width]),axis=0)/255.0 st = time.time() output_image = sess.run(network,feed_dict={net_input:input_image}) run_time = time.time()-st output_image = np.array(output_image[0,:,:,:]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation(output_image, label_values) image_path_no_ext = os.path.splitext(image_path)[0] file_name = utils.filepath_to_name(image_path_no_ext) cv2.imwrite(os.path.join(args.output_folder, "%s.png"%(file_name)),cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
def pred(img): ''' parser = argparse.ArgumentParser() parser.add_argument('--image', type=str, default=None, required=True, help='The image you want to predict on. ') parser.add_argument('--checkpoint_path', type=str, default=None, required=True, help='The path to the latest checkpoint weights for your model.') parser.add_argument('--crop_height', type=int, default=512, help='Height of cropped input image to network') parser.add_argument('--crop_width', type=int, default=512, help='Width of cropped input image to network') parser.add_argument('--model', type=str, default=None, required=True, help='The model you are using') parser.add_argument('--dataset', type=str, default="CamVid", required=False, help='The dataset you are using') ''' print("Testing image " + img) loaded_image = utils.load_image(img) resized_image = cv2.resize(loaded_image, (crop_width, crop_height)) input_image = np.expand_dims( np.float32(resized_image[:crop_height, :crop_width]), axis=0) / 255.0 st = time.time() output_image = sess.run(network, feed_dict={net_input: input_image}) run_time = time.time() - st output_image = np.array(output_image[0, :, :, :]) output_image = helpers.reverse_one_hot(output_image) out_vis_image = helpers.colour_code_segmentation(output_image, label_values) file_name = utils.filepath_to_name(img) res = cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR) cv2.imwrite("%s_pred.png" % (file_name), res) final_out = cv2.addWeighted(resized_image, 1, res, 0.6, 0) cv2.imwrite("%s_pred_overlayed.png" % (file_name), final_out) print("") print("Finished!") print("Wrote image " + "%s_pred.png" % (file_name)) # USAGE # python segment.py --model enet-cityscapes/enet-model.net --classes enet-cityscapes/enet-classes.txt --colors enet-cityscapes/enet-colors.txt --image images/example_01.png # load the class label names CLASSES = open(classes).read().strip().split("\n") # if a colors file was supplied, load it from disk if colors: COLORS = open(colors).read().strip().split("\n") COLORS = [np.array(c.split(",")).astype("int") for c in COLORS] COLORS = np.array(COLORS, dtype="uint8") #print(COLORS) legend = np.zeros(((len(CLASSES) * 15) + 15, 300, 3), dtype="uint8") + 255 # loop over the class names + colors for (i, (className, color)) in enumerate(zip(CLASSES, COLORS)): # draw the class name + color on the legend color = [int(c) for c in color] cv2.putText(legend, ' ' + className, (0, (i * 15) + 14), cv2.FONT_HERSHEY_SIMPLEX, 0.4, tuple(color), 1, cv2.LINE_AA) cv2.rectangle(legend, (100, (i * 15)), (300, (i * 15) + 15), tuple(color), -1) leg = cv2.cvtColor(np.uint8(legend), cv2.COLOR_RGB2BGR) cv2.imwrite("legend.png", leg) return [img, file_name + '_pred_overlayed.png', 'legend.png']