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