diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..496ee2c
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+.DS_Store
\ No newline at end of file
diff --git a/README.md b/README.md
index 2418979..13ce97b 100644
--- a/README.md
+++ b/README.md
@@ -104,6 +104,10 @@ Images are saved to `./RESULT/DATASET/TEST_NAME/`
You can donwload the pretrained weights from here.
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 ChromaGAN: An Adversarial Approach for Picture Colorization:
diff --git a/SAMPLE_IMGS/color/bruschetta.png b/SAMPLE_IMGS/color/bruschetta.png
new file mode 100644
index 0000000..c90fc33
Binary files /dev/null and b/SAMPLE_IMGS/color/bruschetta.png differ
diff --git a/SAMPLE_IMGS/color/fisherman.png b/SAMPLE_IMGS/color/fisherman.png
new file mode 100644
index 0000000..317d5bc
Binary files /dev/null and b/SAMPLE_IMGS/color/fisherman.png differ
diff --git a/SAMPLE_IMGS/color/ocean.png b/SAMPLE_IMGS/color/ocean.png
new file mode 100644
index 0000000..b1f6d77
Binary files /dev/null and b/SAMPLE_IMGS/color/ocean.png differ
diff --git a/SAMPLE_IMGS/grayscale/bruschetta.png b/SAMPLE_IMGS/grayscale/bruschetta.png
new file mode 100644
index 0000000..bbe615f
Binary files /dev/null and b/SAMPLE_IMGS/grayscale/bruschetta.png differ
diff --git a/SAMPLE_IMGS/grayscale/fisherman.png b/SAMPLE_IMGS/grayscale/fisherman.png
new file mode 100644
index 0000000..69d69ba
Binary files /dev/null and b/SAMPLE_IMGS/grayscale/fisherman.png differ
diff --git a/SAMPLE_IMGS/grayscale/ocean.png b/SAMPLE_IMGS/grayscale/ocean.png
new file mode 100644
index 0000000..b38c0c5
Binary files /dev/null and b/SAMPLE_IMGS/grayscale/ocean.png differ
diff --git a/SOURCE/ChromaGANPrint.py b/SOURCE/ChromaGANPrint.py
index 42700ea..3a70cc1 100644
--- a/SOURCE/ChromaGANPrint.py
+++ b/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 .
+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)
diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu
new file mode 100644
index 0000000..ef0a91a
--- /dev/null
+++ b/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"]
diff --git a/docker/Dockerfile.gpu b/docker/Dockerfile.gpu
new file mode 100644
index 0000000..6b69fb6
--- /dev/null
+++ b/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"]
diff --git a/docker/build-and-push.sh b/docker/build-and-push.sh
new file mode 100755
index 0000000..31b9092
--- /dev/null
+++ b/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