Skip to content

TensorFlow implementation of the Efficient Sub-Pixel Convolutional Neural Network

License

Notifications You must be signed in to change notification settings

kweisamx/TensorFlow-ESPCN

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

25 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Implement ESCPN with Tensorflow

Imgur

Dependency

pip

  • Tensorflow
  • Opencv
  • h5py

How to train

python main.py

if you want to see all flag

python main.py -h

How to test

If you don't input a Test image, it will be default image

python main.py --is_train False

then result will put in the result directory

If you want to Test your own iamge

use test_img flag

python main.py --is_train False --test_img Train/t20.bmp

then result image also put in the result directory

Subpixel CNN layer

source

In numpy, we can write this as

def PS(I, r):
  assert len(I.shape) == 3
  assert r>0
  r = int(r)
  O = np.zeros((I.shape[0]*r, I.shape[1]*r, I.shape[2]/(r*2)))
  for x in range(O.shape[0]):
    for y in range(O.shape[1]):
      for c in range(O.shape[2]):
        c += 1
        a = np.floor(x/r).astype("int")
        b = np.floor(y/r).astype("int")
        d = c*r*(y%r) + c*(x%r)
        print a, b, d
        O[x, y, c-1] = I[a, b, d]
  return O

To implement this in Tensorflow we would have to create a custom operator and its equivalent gradient. But after staring for a few minutes in the image depiction of the resulting operation we noticed how to write that using just regular reshape, split and concatenate operations. To understand that note that phase shift simply goes through different channels of the output convolutional map and builds up neighborhoods of r x r pixels. And we can do the same with a few lines of Tensorflow code as:

   def _phase_shift(self, I, r):
       # Helper function with main phase shift operation
       bsize, a, b, c = I.get_shape().as_list()
       X = tf.reshape(I, (self.batch_size, a, b, r, r))
       X = tf.split(X, a, 1)  # a, [bsize, b, r, r]
       X = tf.concat([tf.squeeze(x) for x in X], 2)  # bsize, b, a*r, r
       X = tf.split(X, b, 1)  # b, [bsize, a*r, r]
       X = tf.concat([tf.squeeze(x) for x in X], 2)  # bsize, a*r, b*r
       return tf.reshape(X, (self.batch_size, a*r, b*r, 1))

   def PS(self, X, r):
       # Main OP that you can arbitrarily use in you tensorflow code
       Xc = tf.split(X, 3, 3)
       if self.is_train:
           X = tf.concat([self._phase_shift(x, r) for x in Xc], 3) # Do the concat RGB
       else:
           X = tf.concat([self._phase_shift_test(x, r) for x in Xc], 3) # Do the concat RGB
       return X

Result

  • origin 255 x 255 x 3

Imgur

  • upscaling 3 times, 765 x 765 x 3

Imgur

References

problem

If you meet the problem with opencv when run the program

libSM.so.6: cannot open shared object file: No such file or directory

please install dependency package

sudo apt-get install libsm6
sudo apt-get install libxrender1

About

TensorFlow implementation of the Efficient Sub-Pixel Convolutional Neural Network

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Languages