def run_demo(seq_name, max_training_iters=500, **kwargs): # User Defined parameters gpu_id = kwargs.get('gpu_id', 0) train_model = kwargs.get('train_model', True) side_supervision = kwargs.get('side_supervision', 1) # define paths result_path = os.path.join('DAVIS', 'Results', 'Segmentations', '480p', 'OSVOS', seq_name) parent_path = os.path.join('models', 'OSVOS_parent', 'OSVOS_parent.ckpt-50000') logs_path = os.path.join('models', seq_name) # Train parameters learning_rate = kwargs.get('learning_rate', 1e-8) save_step = kwargs.get('save_step', max_training_iters) display_step = kwargs.get('display_step', 10) # Define Dataset test_frames = sorted( os.listdir(os.path.join('DAVIS', 'JPEGImages', '480p', seq_name))) test_imgs = [ os.path.join('DAVIS', 'JPEGImages', '480p', seq_name, frame) for frame in test_frames ] if train_model: train_imgs = [ os.path.join('DAVIS', 'JPEGImages', '480p', seq_name, '00000.jpg') + ' ' + os.path.join('DAVIS', 'Annotations', '480p', seq_name, '00000.png') ] dataset = Dataset(train_imgs, test_imgs, './', data_aug=True) else: dataset = Dataset(None, test_imgs, './') # Train the network if train_model: with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): global_step = tf.Variable(0, name='global_step', trainable=False) osvos.train_finetune(dataset, parent_path, side_supervision, learning_rate, logs_path, max_training_iters, save_step, global_step, iter_mean_grad=1, ckpt_name=seq_name) # Test the network with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): checkpoint_path = os.path.join( 'models', seq_name, seq_name + '.ckpt-' + str(max_training_iters)) osvos.test(dataset, checkpoint_path, result_path)
def test_model(self): with tf.Graph().as_default(): checkpoint_path = "{}/{}.ckpt-{}".format(self.model_result_path, self.seq_name, self.train_iters) osvos.test(self.dataset, checkpoint_path, self.result_path) pass
def run_osvos_(imgs_dir, labels_dir, result_dir, max_training_iters=500): # User defined parameters gpu_id = 0 train_model = True result_path = result_dir # Train parameters seq_name = 'osvos' parent_path = os.path.join('models', 'OSVOS_parent', 'OSVOS_parent.ckpt-50000') logs_path = os.path.join('models', 'osvos') # Define Dataset test_frames = sorted(os.listdir(imgs_dir)) test_imgs = [os.path.join(imgs_dir, frame) for frame in test_frames] label_frames = sorted(os.listdir(labels_dir)) if train_model: train_imgs = [ os.path.join(imgs_dir, frame[:-4] + '.jpg') + ' ' + os.path.join(labels_dir, frame) for frame in label_frames ] dataset = Dataset(train_imgs, test_imgs, './', data_aug=True) else: dataset = Dataset(None, test_imgs, './') # Train the network if train_model: # More training parameters learning_rate = 1e-8 save_step = max_training_iters side_supervision = 3 display_step = 10 with tf.Graph().as_default(): with tf.device('/cpu:' + str(gpu_id)): global_step = tf.Variable(0, name='global_step', trainable=False) osvos.train_finetune(dataset, parent_path, side_supervision, learning_rate, logs_path, max_training_iters, save_step, display_step, global_step, iter_mean_grad=1, ckpt_name=seq_name) # Test the network with tf.Graph().as_default(): with tf.device('/cpu:' + str(gpu_id)): checkpoint_path = os.path.join( 'models', seq_name, seq_name + '.ckpt-' + str(max_training_iters)) osvos.test(dataset, checkpoint_path, result_path)
def test_parent(seq_name, **kwargs): # result path result_path = os.path.join('DAVIS', 'Results', 'Segmentations', '480p', 'OSVOS-parent', seq_name) # Define Dataset test_frames = sorted(os.listdir(os.path.join('DAVIS', 'JPEGImages', '480p', seq_name))) test_imgs = [os.path.join('DAVIS', 'JPEGImages', '480p', seq_name, frame) for frame in test_frames] dataset = Dataset(None, test_imgs, './') with tf.Graph().as_default(): checkpoint_path = os.path.join('models', 'OSVOS_parent', 'OSVOS_parent.ckpt-50000') osvos.test(dataset, checkpoint_path, result_path)
def start(num_classes, train_path, test_images_path, result_path, parent_path, logs_path, seq_name): gpu_id = 0 train_model = True max_training_iters = 20000 # Define Dataset test_frames = sorted(os.listdir(test_images_path)) test_imgs = [os.path.join(test_images_path, frame) for frame in test_frames] if train_model: dataset = Dataset(train_path, test_imgs, './', data_aug=True) else: dataset = Dataset(None, test_imgs, './') # Train the network if train_model: # More training parameters learning_rate = 1e-8 save_step = max_training_iters side_supervision = 3 display_step = 10 batch_size = 16 with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): global_step = tf.Variable(0, name='global_step', trainable=False) osvos.train_finetune(dataset, num_classes, parent_path, side_supervision, learning_rate, logs_path, max_training_iters, save_step, display_step, global_step, iter_mean_grad=1, ckpt_name=seq_name, batch_size = batch_size) import datetime starttime = datetime.datetime.now() # Test the network with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): checkpoint_path = os.path.join(logs_path, seq_name+'.ckpt-'+str(max_training_iters)) osvos.test(dataset, num_classes, checkpoint_path, result_path) endtime = datetime.datetime.now() print 'Over {0}: Escape time '.format(seq_name) print (endtime - starttime)
# More training parameters learning_rate = 1e-8 save_step = max_training_iters side_supervision = 3 display_step = 10 with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): global_step = tf.Variable(0, name='global_step', trainable=False) osvos.train_finetune(dataset, parent_path, side_supervision, learning_rate, logs_path, max_training_iters, save_step, display_step, global_step, iter_mean_grad=1, ckpt_name=seq_name) # Test the network with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): checkpoint_path = os.path.join('models', seq_name, seq_name+'.ckpt-'+str(max_training_iters)) osvos.test(dataset, checkpoint_path, result_path) # Show results overlay_color = [255, 0, 0] transparency = 0.6 plt.ion() for img_p in test_frames: frame_num = img_p.split('.')[0] img = np.array(Image.open(os.path.join('DAVIS', 'JPEGImages', '480p', seq_name, img_p))) mask = np.array(Image.open(os.path.join(result_path, frame_num+'.png'))) mask = mask/np.max(mask) im_over = np.ndarray(img.shape) im_over[:, :, 0] = (1 - mask) * img[:, :, 0] + mask * (overlay_color[0]*transparency + (1-transparency)*img[:, :, 0]) im_over[:, :, 1] = (1 - mask) * img[:, :, 1] + mask * (overlay_color[1]*transparency + (1-transparency)*img[:, :, 1]) im_over[:, :, 2] = (1 - mask) * img[:, :, 2] + mask * (overlay_color[2]*transparency + (1-transparency)*img[:, :, 2]) plt.imshow(im_over.astype(np.uint8))
def demo(seq_name): # first read in ground truth segmentation and determine how many object are there annatation_0 = os.path.join('..', 'davis-2017', 'data', 'DAVIS', 'Annotations', '480p', seq_name, '00000.png') image_0 = os.path.join('..', 'davis-2017', 'data', 'DAVIS', 'JPEGImages', '480p', seq_name, '00000.jpg') an, _ = io.imread_indexed(annatation_0) N_OBJECT = len(np.unique(an)) - 1 for n_object_ in range(1, N_OBJECT + 1): n_object = str(n_object_) result_path = os.path.join('..', 'davis-2017', 'data', 'DAVIS', 'Results', 'Segmentations', '480p', 'OSVOS2', seq_name, n_object) logs_path = os.path.join('models2', seq_name, n_object) if os.path.exists(os.path.join(logs_path, "done")): continue # Define Dataset test_frames = sorted( os.listdir( os.path.join('..', 'davis-2017', 'data', 'DAVIS', 'JPEGImages', '480p', seq_name))) test_imgs = [ os.path.join('..', 'davis-2017', 'data', 'DAVIS', 'JPEGImages', '480p', seq_name, frame) for frame in test_frames ] if train_model: # we need to first create a new annotation that has only one object base_image = np.zeros_like(an).astype('uint8') base_image[an == n_object_] = n_object_ custom_anno = os.path.join(CUSTOM_ANNOTATION_DIR, seq_name, n_object) if not os.path.exists(custom_anno): os.makedirs(custom_anno) io.imwrite_indexed(os.path.join(custom_anno, "00000.png"), base_image) train_imgs = [ image_0 + ' ' + os.path.join(custom_anno, "00000.png") ] dataset = Dataset(train_imgs, test_imgs, './', data_aug=True) else: dataset = Dataset(None, test_imgs, './') # Train the network if train_model: # More training parameters learning_rate = 1e-8 save_step = max_training_iters side_supervision = 3 display_step = 10 with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): # with tf.device('/cpu:0'): # import pdb; pdb.set_trace() global_step = tf.Variable(0, name='global_step', trainable=False) osvos.train_finetune(dataset, parent_path, side_supervision, learning_rate, logs_path, max_training_iters, save_step, display_step, global_step, iter_mean_grad=1, ckpt_name=seq_name) os.makedirs(os.path.join(logs_path, "done")) # import pdb; pdb.set_trace() # Test the network if os.path.exists(os.path.join(result_path, "done")): continue with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): checkpoint_path = os.path.join( 'models2', seq_name, n_object, seq_name + '.ckpt-' + str(max_training_iters)) osvos.test(dataset, checkpoint_path, result_path) os.makedirs(os.path.join(result_path, "done"))
learning_rate, logs_path, max_training_iters, save_step, display_step, global_step, iter_mean_grad=1, ckpt_name=seq_name, backbone='resnet') # Test the network with tf.Graph().as_default(): with tf.device('/gpu:' + str(gpu_id)): # checkpoint_path = os.path.join('models', seq_name, seq_name+'.ckpt-'+str(max_training_iters)) checkpoint_path = parent_path osvos.test(dataset, checkpoint_path, result_path, backbone='vgg') # Show results overlay_color = [255, 0, 0] transparency = 0.6 plt.ion() for img_p in test_imgs: frame_num = os.path.basename(img_p).split('.')[0] img = np.array(Image.open(img_p)) mask = np.array(Image.open(os.path.join(result_path, frame_num + '.png'))) mask = mask // np.max(mask) mask = 1 - mask im_over = np.ndarray(img.shape) im_over[:, :, 0] = (1 - mask) * img[:, :, 0] + mask * ( overlay_color[0] * transparency + (1 - transparency) * img[:, :, 0]) im_over[:, :, 1] = (1 - mask) * img[:, :, 1] + mask * (