Colorizing Black and White Images with A.I.

git clone https://github.com/piddnad/DDColor.git
cd DDColor

mkdir modelscope

# create a new venv inside the DDColor project
python -m venv venv

venv/bin/pip3 install modelscope

venv/bin/python3 -c "from modelscope.hub.snapshot_download import snapshot_download; snapshot_download('damo/cv_ddcolor_image-colorization', cache_dir='./modelscope')"

import os
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from cog import BasePredictor, Input, Path
from basicsr.archs.ddcolor_arch import DDColor

def color_transfer(reference_bgr: np.ndarray, target_bgr: np.ndarray) -> np.ndarray:
    """Transfer color palette of reference image to target image using Lab mean/std."""
    # Convert BGR images to Lab color space (float32 for precision)
    ref_lab = cv2.cvtColor(reference_bgr, cv2.COLOR_BGR2LAB).astype("float32")
    tgt_lab = cv2.cvtColor(target_bgr, cv2.COLOR_BGR2LAB).astype("float32")
    # Split channels
    L_ref, a_ref, b_ref = cv2.split(ref_lab)
    L_tgt, a_tgt, b_tgt = cv2.split(tgt_lab)
    # Compute mean and std for reference and target (a and b channels)
    a_ref_mean, a_ref_std = a_ref.mean(), a_ref.std()
    b_ref_mean, b_ref_std = b_ref.mean(), b_ref.std()
    a_tgt_mean, a_tgt_std = a_tgt.mean(), a_tgt.std()
    b_tgt_mean, b_tgt_std = b_tgt.mean(), b_tgt.std()
    # Subtract target means, scale by reference/target std, add reference means
    # (Avoid division by zero in case of zero std)
    if a_tgt_std > 1e-6:
        a_tgt = ((a_tgt - a_tgt_mean) * (a_ref_std / a_tgt_std)) + a_ref_mean
    else:
        a_tgt = a_tgt - a_tgt_mean + a_ref_mean
    if b_tgt_std > 1e-6:
        b_tgt = ((b_tgt - b_tgt_mean) * (b_ref_std / b_tgt_std)) + b_ref_mean
    else:
        b_tgt = b_tgt - b_tgt_mean + b_ref_mean
    # Clip values to valid Lab range [0,255] after transfer
    a_tgt = np.clip(a_tgt, 0, 255)
    b_tgt = np.clip(b_tgt, 0, 255)
    # Merge channels back (use original L from target to preserve brightness)
    merged_lab = cv2.merge([L_tgt, a_tgt, b_tgt]).astype("uint8")
    # Convert back to BGR color space
    result_bgr = cv2.cvtColor(merged_lab, cv2.COLOR_LAB2BGR)
    return result_bgr

class ImageColorizationPipeline:
    """Helper pipeline to load DDColor model and process images."""
    def __init__(self, model_path: str, input_size: int = 256, model_size: str = "large"):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # Choose encoder backbone based on model size
        encoder_name = "convnext-t" if model_size == "tiny" else "convnext-l"
        # Initialize DDColor model
        self.model = DDColor(
            encoder_name=encoder_name,
            decoder_name="MultiScaleColorDecoder",
            input_size=[input_size, input_size],
            num_output_channels=2,    # model predicts ab channels
            last_norm="Spectral",
            do_normalize=False,
            num_queries=100,
            num_scales=3,
            dec_layers=9,
        ).to(self.device)
        # Load weights
        state = torch.load(model_path, map_location="cpu")
        # Some checkpoints store weights under 'params' key
        self.model.load_state_dict(state.get("params", state), strict=False)
        self.model.eval()
        self.input_size = input_size  # store for processing

    @torch.no_grad()
    def process(self, img_bgr: np.ndarray) -> np.ndarray:
        """Colorize a BGR image (numpy array) using the loaded DDColor model."""
        # Preserve original resolution L channel
        orig_h, orig_w = img_bgr.shape[:2]
        # Convert to Lab and extract L channel at original size
        img_lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB).astype("float32")
        orig_L = img_lab[:, :, 0:1]  # shape (H, W, 1)
        # Prepare grayscale input at model resolution
        # Resize input to model expected size
        inp = cv2.resize(img_bgr, (self.input_size, self.input_size))
        inp_lab = cv2.cvtColor(inp, cv2.COLOR_BGR2LAB).astype("float32")
        L = inp_lab[:, :, 0:1]
        # Create grayscale Lab image by zeroing AB channels
        gray_lab = np.concatenate([L, np.zeros_like(L), np.zeros_like(L)], axis=2)
        # Convert back to RGB (now it's a gray RGB image of size input_size)
        gray_rgb = cv2.cvtColor(gray_lab.astype("uint8"), cv2.COLOR_LAB2RGB)
        # Prepare tensor and run model
        tensor = torch.from_numpy(gray_rgb.transpose(2, 0, 1)).float().unsqueeze(0).to(self.device)
        out_ab = self.model(tensor)  # output is (1,2,H_out,W_out) in Lab AB
        out_ab = out_ab.cpu().float()
        # Resize output AB to original image size
        out_ab_resized = F.interpolate(out_ab, size=(orig_h, orig_w), mode="bilinear", align_corners=False)
        out_ab_resized = out_ab_resized[0].numpy().transpose(1, 2, 0)  # (H, W, 2)
        # Combine original L and predicted AB, then convert to BGR
        out_lab = np.concatenate([orig_L, out_ab_resized], axis=2).astype("uint8")
        out_bgr = cv2.cvtColor(out_lab, cv2.COLOR_LAB2BGR)
        return out_bgr

class Predictor(BasePredictor):
    def setup(self):
        """Load models into memory for efficient multiple predictions."""
        # Determine paths for large and tiny model weights
        large_model_path = "checkpoints/ddcolor_modelscope.pth"
        tiny_model_path  = "checkpoints/ddcolor_paper_tiny.pth"
        # (Ensure the above files exist. If not, download from ModelScope or HuggingFace as described.)
        # Initialize pipelines for large and tiny models
        self.colorizer_large = ImageColorizationPipeline(model_path=large_model_path, input_size=512, model_size="large")
        self.colorizer_tiny  = ImageColorizationPipeline(model_path=tiny_model_path,  input_size=512, model_size="tiny")

    def predict(
        self,
        image: Path = Input(description="Grayscale input image."),
        reference: Path = Input(description="Optional reference image to guide colors", default=None),
        model_size: str = Input(description="Model size to use (large = highest quality, tiny = faster).",
                                choices=["large", "tiny"], default="large"),
    ) -> Path:
        """Run a single prediction. Colorizes the image, optionally using a reference image for color bias."""
        # Read input image
        img = cv2.imread(str(image))
        if img is None:
            raise ValueError("Failed to load input image")
        # Select model pipeline
        pipeline = self.colorizer_tiny if model_size == "tiny" else self.colorizer_large
        # Colorize the image using DDColor
        output_bgr = pipeline.process(img)
        # If a reference image is provided, transfer its color style to the output
        if reference is not None:
            ref_img = cv2.imread(str(reference))
            if ref_img is None:
                raise ValueError("Failed to load reference image")
            # Only apply transfer if reference is a color image
            if ref_img.shape[2] == 3:
                output_bgr = color_transfer(ref_img, output_bgr)
        # Save result
        out_path = "/tmp/out.png"
        cv2.imwrite(out_path, output_bgr)
        return Path(out_path)