Upload 9 files

.gitattributes

@@ -0,0 +1,33 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Menghan Xia
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,13 @@
+---
+title: colorizator
+emoji: 🐢
+colorFrom: pink
+colorTo: indigo
+sdk: gradio
+sdk_version: 3.9
+app_file: app.py
+pinned: false
+license: openrail
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,91 @@
+import gradio as gr
+import os, requests
+import numpy as np
+from inference import setup_model, colorize_grayscale, predict_anchors
+
+## local |  remote
+RUN_MODE = "remote"
+if RUN_MODE != "local":
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/disco-beta.pth.rar")
+    os.rename("disco-beta.pth.rar", "./checkpoints/disco-beta.pth.rar")
+    ## examples
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/01.jpg")
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/02.jpg")
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/03.jpg")
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/04.jpg")
+
+## step 1: set up model
+device = "cpu"
+checkpt_path = "checkpoints/disco-beta.pth.rar"
+colorizer, colorLabeler = setup_model(checkpt_path, device=device)
+
+def click_colorize(rgb_img, hint_img, n_anchors, is_high_res, is_editable):
+    if hint_img is None:
+        hint_img = rgb_img
+    output = colorize_grayscale(colorizer, colorLabeler, rgb_img, hint_img, n_anchors, True, is_editable, device)
+    output1 = colorize_grayscale(colorizer, colorLabeler, rgb_img, hint_img, n_anchors, False, is_editable, device)
+    return output, output1
+
+def click_predanchors(rgb_img, n_anchors, is_high_res, is_editable):
+    output = predict_anchors(colorizer, colorLabeler, rgb_img, n_anchors, is_high_res, is_editable, device)
+    return output
+
+## step 2: configure interface
+def switch_states(is_checked):
+    if is_checked:
+        return gr.Image.update(visible=True), gr.Button.update(visible=True)
+    else:
+        return gr.Image.update(visible=False), gr.Button.update(visible=False)
+
+demo = gr.Blocks(title="DISCO")
+with demo:
+    gr.Markdown(value="""
+                    **Gradio demo for DISCO: Disentangled Image Colorization via Global Anchors**. Check our [project page](https://menghanxia.github.io/projects/disco.html) 😛.
+                    """)
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                Image_input = gr.Image(type="numpy", label="Input", interactive=True)
+                Image_anchor = gr.Image(type="numpy", label="Anchor", tool="color-sketch", interactive=True, visible=False)
+            with gr.Row():
+                Num_anchor = gr.Number(type="int", value=8, label="Num. of anchors (3~14)")
+                Radio_resolution = gr.Radio(type="index", choices=["Low (256x256)", "High (512x512)"], \
+                                                label="Colorization resolution (Low is more stable)", value="Low (256x256)")
+            with gr.Row():
+                Ckeckbox_editable = gr.Checkbox(default=False, label='Show editable anchors')
+                Button_show_anchor = gr.Button(value="Predict anchors", visible=False)
+            Button_run = gr.Button(value="Colorize")
+        with gr.Column():
+            Image_output = [gr.Image(type="numpy", label="Output").style(height=480), gr.Image(type="numpy", label="Output").style(height=480)]
+
+    Ckeckbox_editable.change(fn=switch_states, inputs=Ckeckbox_editable, outputs=[Image_anchor, Button_show_anchor])
+    Button_show_anchor.click(fn=click_predanchors, inputs=[Image_input, Num_anchor, Radio_resolution, Ckeckbox_editable], outputs=Image_anchor)
+    Button_run.click(fn=click_colorize, inputs=[Image_input, Image_anchor, Num_anchor, Radio_resolution, Ckeckbox_editable], \
+                    outputs=Image_output)
+    
+    ## guiline
+    gr.Markdown(value="""    
+                    🔔**Guideline**
+                    1. Upload your image or select one from the examples.
+                    2. Set up the arguments: "Num. of anchors" and "Colorization resolution".
+                    3. Run the colorization (two modes supported):
+                        - 📀Automatic mode: **Click** "Colorize" to get the automatically colorized output.
+                        - ✏️Editable mode: **Check** ""Show editable anchors"; **Click** "Predict anchors"; **Redraw** the anchor colors (only anchor region will be used); **Click** "Colorize" to get the result.
+                    """)
+    if RUN_MODE != "local":
+        gr.Examples(examples=[
+                    ['01.jpg', 8, "Low (256x256)"],
+                    ['02.jpg', 8, "Low (256x256)"],
+                    ['03.jpg', 8, "Low (256x256)"],
+                    ['04.jpg', 8, "Low (256x256)"],
+                    ], 
+                    inputs=[Image_input,Num_anchor,Radio_resolution], outputs=[Image_output], label="Examples", cache_examples=False)
+    gr.HTML(value="""
+                <p style="text-align:center; color:orange"><a href='https://menghanxia.github.io/projects/disco.html' target='_blank'>DISCO Project Page</a> | <a href='https://github.com/MenghanXia/DisentangledColorization' target='_blank'>Github Repo</a></p>
+                    """)
+
+if RUN_MODE == "local":
+    demo.launch(server_name='9.134.253.83',server_port=7788)
+else:
+    demo.queue(default_enabled=True, status_update_rate=5)
+    demo.launch()

cog.yaml

@@ -0,0 +1,41 @@
+# Configuration for Cog ⚙️
+# Reference: https://github.com/replicate/cog/blob/main/docs/yaml.md
+
+build:
+  # set to true if your model requires a GPU
+  cuda: "10.2"
+  gpu: true
+
+  # a list of ubuntu apt packages to install
+  system_packages:
+    # - "libgl1-mesa-glx"
+    # - "libglib2.0-0"
+    - "libgl1-mesa-dev"
+
+  # python version in the form '3.8' or '3.8.12'
+  python_version: "3.8"
+
+  # a list of packages in the format <package-name>==<version>
+  python_packages:
+    # - "numpy==1.19.4"
+    # - "torch==1.8.0"
+    # - "torchvision==0.9.0"
+    - "numpy==1.23.1"
+    - "torch==1.8.0"
+    - "torchvision==0.9.0"
+    - "opencv-python==4.6.0.66"
+    - "pandas==1.4.3"
+    - "pillow==9.2.0"
+    - "tqdm==4.64.0"
+    - "scikit-image==0.19.3"
+    - "scikit-learn==1.1.2"
+    - "scipy==1.9.1"
+  
+  # commands run after the environment is setup
+  # run:
+    # - "echo env is ready!"
+    # - "echo another command if needed"
+
+# predict.py defines how predictions are run on your model
+predict: "predict.py:Predictor"
+#image: "r8.im/menghanxia/disco"

environment.yml

@@ -0,0 +1,121 @@
+name: DISCO
+channels:
+  - pytorch
+  - defaults
+  - conda-forge
+dependencies:
+  - blas=1.0=mkl
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2022.07.19=h06a4308_0
+  - certifi=2022.6.15=py38h06a4308_0
+  - cudatoolkit=10.2.89=hfd86e86_1
+  - freetype=2.11.0=h70c0345_0
+  - giflib=5.2.1=h7b6447c_0
+  - gmp=6.2.1=h295c915_3
+  - gnutls=3.6.15=he1e5248_0
+  - intel-openmp=2021.4.0=h06a4308_3561
+  - jpeg=9b=h024ee3a_2
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=11.2.0=h1234567_1
+  - libiconv=1.16=h7f8727e_2
+  - libidn2=2.3.2=h7f8727e_0
+  - libpng=1.6.37=hbc83047_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtasn1=4.16.0=h27cfd23_0
+  - libtiff=4.1.0=h2733197_1
+  - libunistring=0.9.10=h27cfd23_0
+  - libuv=1.40.0=h7b6447c_0
+  - libwebp=1.2.0=h89dd481_0
+  - lz4-c=1.9.3=h295c915_1
+  - mkl=2021.4.0=h06a4308_640
+  - mkl-service=2.4.0=py38h7f8727e_0
+  - mkl_fft=1.3.1=py38hd3c417c_0
+  - mkl_random=1.2.2=py38h51133e4_0
+  - ncurses=6.3=h5eee18b_3
+  - nettle=3.7.3=hbbd107a_1
+  - ninja=1.10.2=h06a4308_5
+  - ninja-base=1.10.2=hd09550d_5
+  - numpy=1.23.1=py38h6c91a56_0
+  - numpy-base=1.23.1=py38ha15fc14_0
+  - openh264=2.1.1=h4ff587b_0
+  - openssl=1.1.1q=h7f8727e_0
+  - pillow=9.2.0=py38hace64e9_1
+  - pip=22.1.2=py38h06a4308_0
+  - python=3.8.13=h12debd9_0
+  - readline=8.1.2=h7f8727e_1
+  - setuptools=63.4.1=py38h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - sqlite=3.39.2=h5082296_0
+  - tk=8.6.12=h1ccaba5_0
+  - typing_extensions=4.3.0=py38h06a4308_0
+  - wheel=0.37.1=pyhd3eb1b0_0
+  - xz=5.2.5=h7f8727e_1
+  - zlib=1.2.12=h7f8727e_2
+  - zstd=1.4.9=haebb681_0
+  - ffmpeg=4.3=hf484d3e_0
+  - pytorch=1.8.0=py3.8_cuda10.2_cudnn7.6.5_0
+  - torchaudio=0.8.0=py38
+  - torchvision=0.9.0=py38_cu102
+  - pip:
+    - addict==2.4.0
+    - astunparse==1.6.3
+    - cachetools==4.2.4
+    - charset-normalizer==2.0.7
+    - clang==5.0
+    - cycler==0.11.0
+    - flatbuffers==1.12
+    - fonttools==4.37.1
+    - future==0.18.2
+    - gast==0.4.0
+    - google-auth==2.3.2
+    - google-auth-oauthlib==0.4.6
+    - google-pasta==0.2.0
+    - grpcio==1.41.1
+    - h5py==3.1.0
+    - idna==3.3
+    - imageio==2.21.1
+    - joblib==1.1.0
+    - keras==2.6.0
+    - keras-preprocessing==1.1.2
+    - kiwisolver==1.4.4
+    - lpips==0.1.4
+    - markdown==3.3.4
+    - matplotlib==3.5.3
+    - networkx==2.8.6
+    - oauthlib==3.1.1
+    - opencv-python==4.6.0.66
+    - opt-einsum==3.3.0
+    - packaging==21.3
+    - pandas==1.4.3
+    - protobuf==3.19.0
+    - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
+    - pyparsing==3.0.9
+    - python-dateutil==2.8.2
+    - pytz==2022.2.1
+    - pywavelets==1.3.0
+    - pyyaml==6.0
+    - requests==2.26.0
+    - requests-oauthlib==1.3.0
+    - rsa==4.7.2
+    - scikit-image==0.19.3
+    - scikit-learn==1.1.2
+    - scipy==1.9.1
+    - tensorboard-data-server==0.6.1
+    - tensorboard-plugin-wit==1.8.0
+    - tensorflow-estimator==2.6.0
+    - tensorflow-gpu==2.6.0
+    - termcolor==1.1.0
+    - threadpoolctl==3.1.0
+    - tifffile==2022.8.12
+    - torch==1.8.0
+    - tqdm==4.64.0
+    - urllib3==1.26.7
+    - werkzeug==2.0.2
+    - wrapt==1.12.1
+    - yapf==0.32.0
+prefix: /root/data/programs/anaconda3/envs/DISCO
+

inference.py

@@ -0,0 +1,105 @@
+import os, glob, sys, logging
+import argparse, datetime, time
+import numpy as np
+import cv2
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models import model, basic
+from utils import util
+
+
+def setup_model(checkpt_path, device="cuda"):
+    #print('--------------', torch.cuda.is_available())
+    """Load the model into memory to make running multiple predictions efficient"""
+    colorLabeler = basic.ColorLabel(device=device)
+    colorizer = model.AnchorColorProb(inChannel=1, outChannel=313, enhanced=True, colorLabeler=colorLabeler)
+    colorizer = colorizer.to(device)
+    #checkpt_path = "./checkpoints/disco-beta.pth.rar"
+    assert os.path.exists(checkpt_path), "No checkpoint found!"
+    data_dict = torch.load(checkpt_path, map_location=torch.device('cpu'))
+    colorizer.load_state_dict(data_dict['state_dict'])
+    colorizer.eval()
+    return colorizer, colorLabeler
+
+
+def resize_ab2l(gray_img, lab_imgs, vis=False):
+    H, W = gray_img.shape[:2]
+    reszied_ab = cv2.resize(lab_imgs[:,:,1:], (W,H), interpolation=cv2.INTER_LINEAR)
+    if vis:
+        gray_img = cv2.resize(lab_imgs[:,:,:1], (W,H), interpolation=cv2.INTER_LINEAR)
+        return np.concatenate((gray_img[:,:,np.newaxis], reszied_ab), axis=2)
+    else:
+        return np.concatenate((gray_img, reszied_ab), axis=2)
+
+def prepare_data(rgb_img, target_res):
+    rgb_img = np.array(rgb_img / 255., np.float32)
+    lab_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2LAB)
+    org_grays = (lab_img[:,:,[0]]-50.) / 50.
+    lab_img = cv2.resize(lab_img, target_res, interpolation=cv2.INTER_LINEAR)
+        
+    lab_img = torch.from_numpy(lab_img.transpose((2, 0, 1)))
+    gray_img = (lab_img[0:1,:,:]-50.) / 50.
+    ab_chans = lab_img[1:3,:,:] / 110.
+    input_grays = gray_img.unsqueeze(0)
+    input_colors = ab_chans.unsqueeze(0)
+    return input_grays, input_colors, org_grays
+
+
+def colorize_grayscale(colorizer, color_class, rgb_img, hint_img, n_anchors, is_high_res, is_editable, device="cuda"):
+    n_anchors = int(n_anchors)
+    n_anchors = max(n_anchors, 3)
+    n_anchors = min(n_anchors, 14)
+    target_res = (512,512) if is_high_res else (256,256)
+    input_grays, input_colors, org_grays = prepare_data(rgb_img, target_res)
+    input_grays = input_grays.to(device)
+    input_colors = input_colors.to(device)
+    
+    if is_editable:
+        print('>>>:editable mode')
+        sampled_T = -1
+        _, input_colors, _ = prepare_data(hint_img, target_res)
+        input_colors = input_colors.to(device)
+        pal_logit, ref_logit, enhanced_ab, affinity_map, spix_colors, hint_mask = colorizer(input_grays, \
+                                                                input_colors, n_anchors, sampled_T)
+    else:
+        print('>>>:automatic mode')
+        sampled_T = 0
+        pal_logit, ref_logit, enhanced_ab, affinity_map, spix_colors, hint_mask = colorizer(input_grays, \
+                                                                input_colors, n_anchors, sampled_T)
+
+    pred_labs = torch.cat((input_grays,enhanced_ab), dim=1)
+    lab_imgs = basic.tensor2array(pred_labs).squeeze(axis=0)
+    lab_imgs = resize_ab2l(org_grays, lab_imgs)
+        
+    lab_imgs[:,:,0] = lab_imgs[:,:,0] * 50.0 + 50.0
+    lab_imgs[:,:,1:3] = lab_imgs[:,:,1:3] * 110.0
+    rgb_output = cv2.cvtColor(lab_imgs[:,:,:], cv2.COLOR_LAB2RGB)
+    return (rgb_output*255.0).astype(np.uint8)
+
+
+def predict_anchors(colorizer, color_class, rgb_img, n_anchors, is_high_res, is_editable, device="cuda"):
+    n_anchors = int(n_anchors)
+    n_anchors = max(n_anchors, 3)
+    n_anchors = min(n_anchors, 14)
+    target_res = (512,512) if is_high_res else (256,256)
+    input_grays, input_colors, org_grays = prepare_data(rgb_img, target_res)
+    input_grays = input_grays.to(device)
+    input_colors = input_colors.to(device)
+                
+    sampled_T, sp_size = 0, 16
+    pal_logit, ref_logit, enhanced_ab, affinity_map, spix_colors, hint_mask = colorizer(input_grays, \
+                                                            input_colors, n_anchors, sampled_T)
+    pred_probs = pal_logit
+    guided_colors = color_class.decode_ind2ab(ref_logit, T=0)
+    guided_colors = basic.upfeat(guided_colors, affinity_map, sp_size, sp_size)
+    anchor_masks = basic.upfeat(hint_mask, affinity_map, sp_size, sp_size)
+    marked_labs = basic.mark_color_hints(input_grays, guided_colors, anchor_masks, base_ABs=None)
+    lab_imgs = basic.tensor2array(marked_labs).squeeze(axis=0)
+    lab_imgs = resize_ab2l(org_grays, lab_imgs, vis=True)
+        
+    lab_imgs[:,:,0] = lab_imgs[:,:,0] * 50.0 + 50.0
+    lab_imgs[:,:,1:3] = lab_imgs[:,:,1:3] * 110.0
+    rgb_output = cv2.cvtColor(lab_imgs[:,:,:], cv2.COLOR_LAB2RGB)
+    return (rgb_output*255.0).astype(np.uint8)

predict.py

@@ -0,0 +1,104 @@
+# Prediction interface for Cog ⚙️
+# https://github.com/replicate/cog/blob/main/docs/python.md
+
+from cog import BasePredictor, Input, Path
+import tempfile
+import os, glob
+import numpy as np
+import cv2
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models import model, basic
+from utils import util
+
+class Predictor(BasePredictor):
+    def setup(self):
+        seed = 130
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        #print('--------------', torch.cuda.is_available())
+        """Load the model into memory to make running multiple predictions efficient"""
+        self.colorizer = model.AnchorColorProb(inChannel=1, outChannel=313, enhanced=True)
+        self.colorizer = self.colorizer.cuda()
+        checkpt_path = "./checkpoints/disco-beta.pth.rar"
+        assert os.path.exists(checkpt_path)
+        data_dict = torch.load(checkpt_path, map_location=torch.device('cpu'))
+        self.colorizer.load_state_dict(data_dict['state_dict'])
+        self.colorizer.eval()
+        self.color_class = basic.ColorLabel(lambda_=0.5, device='cuda')
+    
+    def resize_ab2l(self, gray_img, lab_imgs):
+        H, W = gray_img.shape[:2]
+        reszied_ab = cv2.resize(lab_imgs[:,:,1:], (W,H), interpolation=cv2.INTER_LINEAR)
+        return np.concatenate((gray_img, reszied_ab), axis=2)
+    
+    def predict(
+        self,
+        image: Path = Input(description="input image. Output will be one or multiple colorized images."),
+        n_anchors: int = Input(
+            description="number of color anchors", ge=3, le=14, default=8
+        ),
+        multi_result: bool = Input(
+            description="to generate diverse results", default=False
+        ),
+        vis_anchors: bool = Input(
+            description="to visualize the anchor locations", default=False
+        )
+    ) -> Path:
+        """Run a single prediction on the model"""
+        bgr_img = cv2.imread(str(image), cv2.IMREAD_COLOR)
+        rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
+        rgb_img = np.array(rgb_img / 255., np.float32)
+        lab_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2LAB)
+        org_grays = (lab_img[:,:,[0]]-50.) / 50.
+        lab_img = cv2.resize(lab_img, (256,256), interpolation=cv2.INTER_LINEAR)
+        
+        lab_img = torch.from_numpy(lab_img.transpose((2, 0, 1)))
+        gray_img = (lab_img[0:1,:,:]-50.) / 50.
+        ab_chans = lab_img[1:3,:,:] / 110.
+        input_grays = gray_img.unsqueeze(0)
+        input_colors = ab_chans.unsqueeze(0)
+        input_grays = input_grays.cuda(non_blocking=True)
+        input_colors = input_colors.cuda(non_blocking=True)
+                
+        sampled_T = 2 if multi_result else 0
+        pal_logit, ref_logit, enhanced_ab, affinity_map, spix_colors, hint_mask = self.colorizer(input_grays, \
+                                                            input_colors, n_anchors, True, sampled_T)
+        pred_probs = pal_logit
+        guided_colors = self.color_class.decode_ind2ab(ref_logit, T=0)
+        sp_size = 16
+        guided_colors = basic.upfeat(guided_colors, affinity_map, sp_size, sp_size)
+        res_list = []
+        if multi_result:
+            for no in range(3):
+                pred_labs = torch.cat((input_grays,enhanced_ab[no:no+1,:,:,:]), dim=1)
+                lab_imgs = basic.tensor2array(pred_labs).squeeze(axis=0)
+                lab_imgs = self.resize_ab2l(org_grays, lab_imgs)
+                #util.save_normLabs_from_batch(lab_imgs, save_dir, [file_name], -1, suffix='c%d'%no)
+                res_list.append(lab_imgs)
+        else:
+            pred_labs = torch.cat((input_grays,enhanced_ab), dim=1)
+            lab_imgs = basic.tensor2array(pred_labs).squeeze(axis=0)
+            lab_imgs = self.resize_ab2l(org_grays, lab_imgs)
+            #util.save_normLabs_from_batch(lab_imgs, save_dir, [file_name], -1)#, suffix='enhanced')
+            res_list.append(lab_imgs)
+        
+        if vis_anchors:
+            ## visualize anchor locations
+            anchor_masks = basic.upfeat(hint_mask, affinity_map, sp_size, sp_size)
+            marked_labs = basic.mark_color_hints(input_grays, enhanced_ab, anchor_masks, base_ABs=enhanced_ab)
+            hint_imgs = basic.tensor2array(marked_labs).squeeze(axis=0)
+            hint_imgs = self.resize_ab2l(org_grays, hint_imgs)
+            #util.save_normLabs_from_batch(hint_imgs, save_dir, [file_name], -1, suffix='anchors')
+            res_list.append(hint_imgs)
+        
+        output = cv2.vconcat(res_list)
+        output[:,:,0] = output[:,:,0] * 50.0 + 50.0
+        output[:,:,1:3] = output[:,:,1:3] * 110.0
+        rgb_output = cv2.cvtColor(output[:,:,:], cv2.COLOR_LAB2BGR)      
+        out_path = Path(tempfile.mkdtemp()) / "out.png"
+        cv2.imwrite(str(out_path), (rgb_output*255.0).astype(np.uint8))
+        return out_path

requirements.txt

@@ -0,0 +1,17 @@
+addict
+future
+numpy
+opencv-python
+pandas
+Pillow
+pyyaml
+requests
+scikit-image
+scikit-learn
+scipy
+torch>=1.8.0
+torchvision
+tensorboardx>=2.4
+tqdm
+yapf
+lpips