Docker image for inference

.gitignore

@@ -0,0 +1 @@
+.DS_Store

README.md

@@ -104,6 +104,10 @@ Images are saved to `./RESULT/DATASET/TEST_NAME/`
 You can donwload the pretrained weights from <a href="https://drive.google.com/drive/folders/12s4rbLmnjW4e8MmESbfRStGbrjOrahlW?usp=sharing">here</a>.
 In order to test the network you should use the file called ` my_model_colorization.h5. 
 
+## Run inferences using Docker
+
+You can colorize arbitrary grayscale images using a pre-packaged Docker image from the Replicate registry: https://beta.replicate.ai/pvitoria/ChromaGAN
+
 ## Citation
 If you use this code for your research, please cite our paper <a href="https://arxiv.org/pdf/1907.09837.pdf"> ChromaGAN: An Adversarial Approach for Picture Colorization</a>:
 

SOURCE/ChromaGANPrint.py

@@ -11,6 +11,8 @@
 #You should have received a copy of the GNU Affero General Public License
 #along with this program. If not, see <http://www.gnu.org/licenses/>.
 
+import sys
+import argparse
 import os
 import tensorflow as tf
 import config as config
@@ -43,7 +45,7 @@ def reconstruct_no(batchX, predictedY):
     return result
 
 
-def sample_images():
+def sample_images(concatenate=True):
     avg_cost = 0
     avg_cost2 = 0
     avg_cost3 = 0
@@ -61,7 +63,11 @@ def sample_images():
     print("total number of batches to colorize " + str(total_batch))
     for b in range(total_batch):
             #batchX, batchY,  filelist  = test_data.generate_batch()
-            batchX, batchY,  filelist, original, labimg_oritList = test_data.generate_batch()
+            try:
+                batchX, batchY,  filelist, original, labimg_oritList = test_data.generate_batch()
+            except Exception as e:
+                sys.stderr.write("Failed to generate batch: {}\n".format(e))
+                continue
             predY, _  = colorizationModel.predict(np.tile(batchX,[1,1,1,3]))
             predictVGG =VGG_modelF.predict(np.tile(batchX,[1,1,1,3]))
             loss = colorizationModel.evaluate(np.tile(batchX,[1,1,1,3]), [batchY, predictVGG], verbose=0)
@@ -82,7 +88,12 @@ def sample_images():
                 labimg_ori =np.expand_dims(labimg_oritList[i],axis=2)
                 predResult= reconstruct_no(deprocess(labimg_ori), predictedAB)
                 save_path = os.path.join(config.OUT_DIR, "{:4.8f}_".format(psnr)+filelist[i][:-4] +"psnr_reconstructed.jpg" )
-                cv2.imwrite(save_path, np.concatenate((predResult, originalResult)))
+                if concatenate:
+                    result_img = np.concatenate((predResult, originalResult))
+                else:
+                    result_img = predResult
+                if not cv2.imwrite(save_path, result_img):
+                    print("Failed to save " + save_path)
                 print("Batch " + str(b)+"/"+str(total_batch))
                 print(psnr)
 
@@ -95,4 +106,7 @@ def sample_images():
 
 
 if __name__ == '__main__':
-    sample_images()
+    parser = argparse.ArgumentParser(description="ChromaGAN colorization")
+    parser.add_argument("--no-concatenate", action="store_true")
+    args = parser.parse_args()
+    sample_images(not args.no_concatenate)

docker/Dockerfile.cpu

@@ -0,0 +1,30 @@
+FROM python:3.7-buster
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && \
+        apt-get install -y \
+        libsm6 \
+        libxext6 \
+        libxrender1
+
+RUN pip install gdown
+RUN gdown -O /tmp/my_model_colorization.h5 'https://drive.google.com/uc?id=12hyrk6A3bcpAFIBNW0sfeEevJ9daUTGr'
+
+RUN pip install \
+        opencv-python==4.1.0.25 \
+        numpy==1.15.4 \
+        tensorflow==1.13.1 \
+        keras==2.2.4 \
+        h5py==2.10.0
+
+RUN mkdir -p /code/MODEL/imagenet && mv /tmp/my_model_colorization.h5 /code/MODEL/imagenet/modelPretrained.h5
+
+# pre-cache vgg16
+RUN python -c "from keras import applications; applications.vgg16.VGG16(weights='imagenet', include_top=True)"
+
+COPY . /code
+WORKDIR /code/SOURCE
+RUN sed -i "s/BATCH_SIZE = 10/BATCH_SIZE = 1/" config.py
+
+ENTRYPOINT ["python", "ChromaGANPrint.py", "--no-concatenate"]

docker/Dockerfile.gpu

@@ -0,0 +1,35 @@
+FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && \
+        apt-get install -y software-properties-common && \
+        add-apt-repository ppa:deadsnakes/ppa && \
+        apt-get update && \
+        apt-get install -y \
+        libsm6 \
+        libxext6 \
+        libxrender1 \
+        python3.7 \
+        python3-pip
+
+RUN pip3 install gdown
+RUN gdown -O /tmp/my_model_colorization.h5 'https://drive.google.com/uc?id=12hyrk6A3bcpAFIBNW0sfeEevJ9daUTGr'
+
+RUN pip3 install \
+        opencv-python==4.1.0.25 \
+        numpy==1.15.4 \
+        tensorflow-gpu==1.13.1 \
+        keras==2.2.4 \
+        h5py==2.10.0
+
+RUN mkdir -p /code/MODEL/imagenet && mv /tmp/my_model_colorization.h5 /code/MODEL/imagenet/modelPretrained.h5
+
+# pre-cache vgg16
+RUN python3 -c "from keras import applications; applications.vgg16.VGG16(weights='imagenet', include_top=True)"
+
+COPY . /code
+WORKDIR /code/SOURCE
+RUN sed -i "s/BATCH_SIZE = 10/BATCH_SIZE = 1/" config.py
+
+ENTRYPOINT ["python3", "ChromaGANPrint.py", "--no-concatenate"]

docker/build-and-push.sh

@@ -0,0 +1,42 @@
+#!/bin/bash -eux
+
+# Run this script from the repo's root folder
+#
+# $ ./docker/build-and-push.sh
+
+# 1. Build Docker images for CPU and GPU
+
+image="us-docker.pkg.dev/replicate/pvitoria/chromagan"
+cpu_tag="$image:cpu"
+gpu_tag="$image:gpu"
+docker build -f docker/Dockerfile.cpu --tag "$cpu_tag" .
+docker build -f docker/Dockerfile.gpu --tag "$gpu_tag" .
+
+# 2. Test Docker images
+
+test_input_folder=/tmp/test-chromagan/input
+mkdir -p $test_input_folder
+cp SAMPLE_IMGS/grayscale/bruschetta.png $test_input_folder/
+test_output_folder=/tmp/test-chromagan/output
+
+docker run -it \
+        -v $test_input_folder:/code/DATASET/imagenet/test \
+        -v $test_output_folder/cpu:/code/RESULT/imagenet \
+        $cpu_tag
+
+[ -f $test_output_folder/cpu/*_bruschettapsnr_reconstructed.jpg ] || exit 1
+
+docker run -it \
+        -v $test_input_folder:/code/DATASET/imagenet/test \
+        -v $test_output_folder/gpu:/code/RESULT/imagenet \
+        $gpu_tag
+
+[ -f $test_output_folder/gpu/*_bruschettapsnr_reconstructed.jpg ] || exit 1
+
+sudo rm -rf "$test_input_folder"
+sudo rm -rf "$test_output_folder"
+
+# 3. Push the Docker images
+
+docker push $cpu_tag
+docker push $gpu_tag