Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updated to Python 3, tensorflow 1.15.2-gpu #1

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
FROM tensorflow/tensorflow:1.15.2-gpu
WORKDIR /wasr
RUN set -euo pipefail && \
export DEBIAN_FRONTEND=noninteractive && \
apt-get update && \
apt-get -y upgrade && \
apt-get -y install --no-install-recommends python-opencv && \
pip3 install opencv-python scipy pillow && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
7 changes: 7 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,12 @@
# Water-obstacle Separation and Refinement network (WaSR)

## Running in a Docker image with NVIDIA GPU support (requires Docker 19.03+)
```
docker build -t wasr .
docker run -it --rm --gpus all -v $(pwd)/:/wasr wasr:latest bash
python wasr_inference_noimu_general.py --img-path example_1.jpg
```

## No-IMU Version
This architecture does not incorporate IMU information. The ARM and FFM modules are used to fuse encoder and decoder features.
Novel water-separation loss is applied early in the encoder (res4 block) to force-separate water pixels from obstacle pixels.
Expand Down
28 changes: 19 additions & 9 deletions kaffe/tensorflow/network.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,8 @@ def feed(self, *args):
assert len(args) != 0
self.terminals = []
for fed_layer in args:
if isinstance(fed_layer, basestring):
#if isinstance(fed_layer, basestring): # py 2.7
if isinstance(fed_layer, str):
try:
#print('Layer ' + fed_layer + ' shape')
#print(self.layers[fed_layer].shape)
Expand Down Expand Up @@ -110,7 +111,8 @@ def validate_padding(self, padding):

@layer
def attention_refinment_module(self, input, name):
global_pool = tf.reduce_mean(input, [1, 2], keep_dims=True)
#global_pool = tf.reduce_mean(input, [1, 2], keep_dims=True) # tf 1.2
global_pool = tf.reduce_mean(input, [1, 2], keepdims=True)
#conv_1 = keras_ly.Conv2D(2048, [1, 1], padding='SAME', name=name+'_conv1')(global_pool)
conv_1 = keras_ly.Conv2D(input.get_shape()[3], [1, 1], padding='SAME', name=name+'_conv1')(global_pool)
sigmoid = tf.sigmoid(conv_1, name=name+'_sigmoid')
Expand All @@ -120,15 +122,17 @@ def attention_refinment_module(self, input, name):

@layer
def attention_refinment_module_new(self, input, name, last_arm=False):
global_pool = tf.reduce_mean(input, [1, 2], keep_dims=True)
#global_pool = tf.reduce_mean(input, [1, 2], keep_dims=True) # tf 1.2
global_pool = tf.reduce_mean(input, [1, 2], keepdims=True)
conv_1 = keras_ly.Conv2D(input.get_shape()[3], [1, 1], padding='SAME', name=name+'_conv1')(global_pool)
with tf.variable_scope(name+'_conv1_bn') as scope:
conv_1_bn = slim.batch_norm(conv_1, fused=True, scope=scope)
sigmoid = tf.sigmoid(conv_1_bn, name=name+'_sigmoid')
mul_out = tf.multiply(input, sigmoid, name=name+'_multiply')

if last_arm:
glob_red = tf.reduce_mean(mul_out, [1, 2], keep_dims=True)
#glob_red = tf.reduce_mean(mul_out, [1, 2], keep_dims=True) # tf 1.2
glob_red = tf.reduce_mean(mul_out, [1, 2], keepdims=True)
out_scale = tf.multiply(glob_red, mul_out)
print('last arm shape')
print(input.shape)
Expand All @@ -147,7 +151,8 @@ def feature_fusion_module(self, input, name):
concat_1 = tf.concat(axis=3, values=[input_big, up_sampled_input], name=name+'_concat')
conv_1 = keras_ly.Conv2D(1024, [3, 3], padding='SAME', name=name+'_conv1')(concat_1)

global_pool = tf.reduce_mean(conv_1, [1, 2], keep_dims=True)
#global_pool = tf.reduce_mean(conv_1, [1, 2], keep_dims=True) # tf 1.2
global_pool = tf.reduce_mean(conv_1, [1, 2], keepdims=True)
conv_2 = keras_ly.Conv2D(1024, [1, 1], padding='SAME', name=name+'_conv2')(global_pool)
conv_3 = keras_ly.Conv2D(1024, [1, 1], padding='SAME', name=name+'_conv3')(conv_2)
sigmoid = tf.sigmoid(conv_3, name=name+'_sigmoid')
Expand Down Expand Up @@ -194,7 +199,8 @@ def multiply_two_tensors(self, input, name):

@layer
def global_pool(self, input, name, axis):
return tf.reduce_mean(input, axis=axis, keep_dims=True, name=name)
#return tf.reduce_mean(input, axis=axis, keep_dims=True, name=name) # tf 1.2
return tf.reduce_mean(input, axis=axis, keepdims=True, name=name)


@layer
Expand All @@ -221,7 +227,8 @@ def conv(self,
# Convolution for a given input and kernel
convolve = lambda i, k: tf.nn.conv2d(i, k, [1, s_h, s_w, 1], padding=padding)
with tf.variable_scope(name) as scope:
kernel = self.make_var('weights', shape=[k_h, k_w, c_i / group, c_o])
#kernel = self.make_var('weights', shape=[k_h, k_w, c_i / group, c_o]) # tf 1.2
kernel = self.make_var('weights', shape=[k_h, k_w, c_i // group, c_o])
if group == 1:
# This is the common-case. Convolve the input without any further complications.
output = convolve(input, kernel)
Expand Down Expand Up @@ -263,7 +270,9 @@ def atrous_conv(self,
# Convolution for a given input and kernel
convolve = lambda i, k: tf.nn.atrous_conv2d(i, k, dilation, padding=padding)
with tf.variable_scope(name) as scope:
kernel = self.make_var('weights', shape=[k_h, k_w, c_i / group, c_o])
#kernel = self.make_var('weights', shape=[k_h, k_w, c_i / group, c_o]) # tf 1.2
kernel = self.make_var('weights', shape=[k_h, k_w, c_i // group, c_o])

if group == 1:
# This is the common-case. Convolve the input without any further complications.
output = convolve(input, kernel)
Expand Down Expand Up @@ -411,4 +420,5 @@ def se_fc(self, input, num_out, name, relu=True):

@layer
def global_pool_layer(self, input, axis, name, keep_dims):
return tf.reduce_mean(input, axis, keep_dims=keep_dims, name=name)
#return tf.reduce_mean(input, axis, keep_dims=keep_dims, name=name) # tf 1.2
return tf.reduce_mean(input, axis, keepdims=keep_dims, name=name)
12 changes: 8 additions & 4 deletions wasr_inference_noimu_general.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,8 @@ def load(saver, sess, ckpt_path):
print("Restored model parameters from {}".format(ckpt_path))

def main():
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
# uncomment/set to an invalid value to run on the CPU instead of the GPU
#os.environ['CUDA_VISIBLE_DEVICES'] = '1'

"""Create the model and start the evaluation process."""
args = get_arguments()
Expand All @@ -91,7 +92,8 @@ def main():
img -= IMG_MEAN

# Expand first dimension
img = tf.expand_dims(img, dim=0)
#img = tf.expand_dims(img, dim=0) # tf 1.2
img = tf.expand_dims(img, axis=0)

with tf.variable_scope('', reuse=False):
net = wasr_NOIMU2({'data': img}, is_training=False, num_classes=args.num_classes)
Expand All @@ -104,8 +106,10 @@ def main():

# Upsample image to the original resolution
raw_output = tf.image.resize_bilinear(raw_output, tf.shape(img)[1:3, ])
raw_output = tf.argmax(raw_output, dimension=3)
pred = tf.expand_dims(raw_output, dim=3)
#raw_output = tf.argmax(raw_output, dimension=3) # tf 1.2
#pred = tf.expand_dims(raw_output, dim=3) # tf 1.2
raw_output = tf.argmax(raw_output, axis=3)
pred = tf.expand_dims(raw_output, axis=3)

# Set up TF session and initialize variables.
config = tf.ConfigProto()
Expand Down