def start(num_classes, train_path, valid_path, parent_path, logs_path, seq_name):
max_training_iters = 3000
# Define BLDataset
dataset = Dataset(train_path, None, './', data_aug=False)
valid_dataset = Dataset(valid_path, None, './', data_aug=False)
# Train the network
#learning_rate = 2e-4
decay_steps = 500
save_step = max_training_iters
side_supervision = 3
display_step = 10
batch_size = 32
starttime = datetime.now()
with tf.Graph().as_default():
with tf.device('/gpu:' + str(gpu_id)):
global_step = tf.Variable(0, name='global_step', trainable=False)
learning_rate = tf.train.exponential_decay(2e-4, global_step,
decay_steps=decay_steps,
decay_rate=0.9)
#staircase=True)
osvos.train_finetune(dataset, valid_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)
endtime = datetime.now()
print 'Over {0}: Escape time '.format(seq_name)
print (endtime - starttime)
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 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 train_model(self,
learning_rate=1e-8,
side_supervision=3,
display_step=10):
if self.is_train:
with tf.Graph().as_default():
global_step = tf.Variable(0,
name='global_step',
trainable=False)
osvos.train_finetune(self.dataset,
self.parent_model,
side_supervision,
learning_rate,
self.model_result_path,
self.train_iters,
self.train_iters,
display_step,
global_step,
iter_mean_grad=1,
ckpt_name=self.seq_name)
pass
pass
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)
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:
# 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')))
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"))
