forked from scaelles/OSVOS-TensorFlow
-
Notifications
You must be signed in to change notification settings - Fork 0
/
osvos_parent_demo.py
67 lines (60 loc) · 2.85 KB
/
osvos_parent_demo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
"""
Sergi Caelles ([email protected])
This file is part of the OSVOS paper presented in:
Sergi Caelles, Kevis-Kokitsi Maninis, Jordi Pont-Tuset, Laura Leal-Taixe, Daniel Cremers, Luc Van Gool
One-Shot Video Object Segmentation
CVPR 2017
Please consider citing the paper if you use this code.
"""
import os
import sys
import tensorflow as tf
slim = tf.contrib.slim
# Import OSVOS files
root_folder = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.abspath(root_folder))
import osvos
from dataset import Dataset
# User defined parameters
gpu_id = 0
# Training parameters
imagenet_ckpt = 'models/vgg_16.ckpt'
logs_path = os.path.join(root_folder, 'models', 'OSVOS_parent')
store_memory = True
data_aug = True
iter_mean_grad = 10
max_training_iters_1 = 15000
max_training_iters_2 = 30000
max_training_iters_3 = 50000
save_step = 5000
test_image = None
display_step = 100
ini_learning_rate = 1e-8
boundaries = [10000, 15000, 25000, 30000, 40000]
values = [ini_learning_rate, ini_learning_rate * 0.1, ini_learning_rate, ini_learning_rate * 0.1, ini_learning_rate,
ini_learning_rate * 0.1]
# Define Dataset
train_file = 'train_parent.txt'
dataset = Dataset(train_file, None, './DAVIS', store_memory=store_memory, data_aug=data_aug)
# Train the network
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.piecewise_constant(global_step, boundaries, values)
osvos.train_parent(dataset, imagenet_ckpt, 1, learning_rate, logs_path, max_training_iters_1, save_step,
display_step, global_step, iter_mean_grad=iter_mean_grad, test_image_path=test_image,
ckpt_name='OSVOS_parent')
with tf.Graph().as_default():
with tf.device('/gpu:' + str(gpu_id)):
global_step = tf.Variable(max_training_iters_1, name='global_step', trainable=False)
learning_rate = tf.train.piecewise_constant(global_step, boundaries, values)
osvos.train_parent(dataset, imagenet_ckpt, 2, learning_rate, logs_path, max_training_iters_2, save_step,
display_step, global_step, iter_mean_grad=iter_mean_grad, resume_training=True,
test_image_path=test_image, ckpt_name='OSVOS_parent')
with tf.Graph().as_default():
with tf.device('/gpu:' + str(gpu_id)):
global_step = tf.Variable(max_training_iters_2, name='global_step', trainable=False)
learning_rate = tf.train.piecewise_constant(global_step, boundaries, values)
osvos.train_parent(dataset, imagenet_ckpt, 3, learning_rate, logs_path, max_training_iters_3, save_step,
display_step, global_step, iter_mean_grad=iter_mean_grad, resume_training=True,
test_image_path=test_image, ckpt_name='OSVOS_parent')