SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zheng Chang#1 Shuchen Weng#2,3 Peixuan Zhang $^{1}$  Yu Li $^{4}$  Si Li $^{*1}$  Boxin Shi $^{2,3}$
$^{1}$  School of Artificial Intelligence, Beijing University of Posts and Telecommunications
$^{2}$  National Key Laboratory for Multimedia Information Processing, School of Computer Science, Peking University
$^{3}$  National Engineering Research Center of Visual Technology, School of Computer Science, Peking University
$^{4}$  International Digital Economy Academy
{zhengchang98,pxzhang,lius}@bupt.edu.cn
{shuchenweng, shiboxin}@pku.edu.cn liyu@idea.edu.cn
Language-based colorization produces plausible and visually pleasing colors under the guidance of user-friendly natural language descriptions. Previous methods implicitly assume that users provide comprehensive color descriptions for most of the objects in the image, which leads to suboptimal performance. In this paper, we propose a unified model to perform language-based colorization with any-level descriptions. We leverage the pretrained cross-modality generative model for its robust language understanding and rich color priors to handle the inherent ambiguity of any-level descriptions. We further design modules to align with input conditions to preserve local spatial structures and prevent the ghosting effect. With the proposed novel sampling strategy, our model achieves instance-aware colorization in diverse and complex scenarios. Extensive experimental results demonstrate our advantages of effectively handling any-level descriptions and outperforming both language-based and automatic colorization methods. The code and pretrained models are available at: https://github.com/changzheng123/L-CAD.
Image colorization is a challenging task that involves converting grayscale images into plausible and visually pleasing colorful ones. Language-based colorization approaches (e.g., [5, 6, 27]) use natural language descriptions as guidance to produce more controllable colorized images. These methods cater to users' specific requests, enabling them to provide more specific and nuanced color preferences. As automatic colorization (e.g., [39, 44, 52]) often encounters color ambiguity for common objects (e.g., flower colors), language-based colorization has shown promising results in producing high-quality and customizable colored images with user-friendly input.
While language-based colorization methods have improved the consistency between language descriptions and colorization results through feature fusion [7, 27], decoupling color-object space [5, 45], and aggregating similar patches [6], they implicitly assume that users provide comprehensive color descriptions for most of the objects in the image. This assumption often leads to suboptimal performance, especially for objects without corresponding color descriptions. Furthermore, we have observed that users commonly assign colors only to their objects of interest.
Equal contribution. * Corresponding author
Figure 1: Colorization with any-level descriptions, comparing language-based colorization methods (e.g., L-CoDe [45] and L-CoIns [6]) and automatic colorization methods (e.g., DISCO [47] and  $\mathrm{CT^2}$  [44]). First row: For complete-level descriptions, automatic colorization methods is unable to colorize images as requested by the user. Second row: For partial-level descriptions, automatic colorization methods fail to colorize mentioned objects with their corresponding colors, while language-based colorization methods struggle to present appropriate colors for unmentioned objects. Third row: For scarce-level descriptions, language-based colorization methods show suboptimal performance without meaningful color hints. In contrast, our method is able to colorize images with any-level descriptions.
To colorize images with descriptions of diverse levels of details, we propose a unified model that adaptively understands any-level descriptions and operates as follows: (i) For complete-level descriptions that cover major objects, the model precisely colorizes the image based on user requests (Fig. 1 first row, comprehensive descriptions for four cups of wines); (ii) for partial-level descriptions that focus on objects of interest, it colorizes unmentioned objects by cuing from image semantics when concentrating on their objects of interest (Fig. 1 second row, selective descriptions for only the jar and flowers); and (iii) for scarce-level descriptions that lack meaningful color information, it switches to the automatic colorization model (Fig. 1 third row, omitted descriptions for the pizza and restaurant).
To achieve the above goal, we present L-CAD to perform Language-based Colorization with Any-level Descriptions. Given the inherent ambiguity in the number of objects mentioned in any-level descriptions, we leverage the pretrained cross-modality generative model (i.e., Stable Diffusion [34]) to utilize its robust language understanding for mentioned objects and rich color priors for unmentioned ones. However, as the generative model is not specifically designed for colorization, it faces challenges in spatial alignment with grayscale inputs. To address this, we develop a luminance-guided compression module that aligns with grayscale images in the pixel space, preserving local spatial structures. We also design a channel-extended convolution operator to prevent the ghosting effect by aligning with descriptions in the latent space. Additionally, to handle descriptions with diverse levels of details and complex scenarios, we develop an instance-aware sampling strategy that roughly estimates object contours in the latent space and assigns color features to their corresponding regions. Our approach is not limited by the number of objects mentioned in descriptions, enabling effective handling of any-level descriptions.
Our contributions could be summarized as follows:
- We propose a unified model that adaptively understands any-level descriptions, achieving state-of-the-art performance in both automatic and language-based colorization methods.
- We develop novel modules to preserve local spatial structures and prevent the ghosting effect by aligning with input conditions in both the pixel space and the latent space.
- We present the instance-aware sampling strategy to correctly assign colors to corresponding objects, enabling effective colorization in diverse and complex scenarios.
2.1 Language-based colorization
Language-based colorization aims to leverage flexible and user-friendly language descriptions to produce visually pleasing results with descriptions that contain objects and their corresponding colors. Manjunatha et al. [27] introduce the FILM pipeline, which employs feature-wise affine transformations to incorporate language descriptions in the grayscale image colorization process. Similarly, Chen et al. [7] utilize a recurrent attentive model to spatially fuse image and language features. To gain a deeper understanding of grayscale images, Xie et al. [49] use additional segmentation masks and introduce semantic segmentation as a side task. L-CoDe [45] and L-CoDer [5] rely on additionally annotated correspondences between object words and color words for guidance, enabling the model to assign specific color words to image regions. To achieve instance-aware colorization, L-CoIns [6] presents the grouping mechanism to automatically aggregate similar image patches. In this paper, we further explore the effective approach to enhance the robustness of language-based colorization, e.g., assigning vivid colors for objects unmentioned in the description.
2.2 Automatic colorization
Automatic colorization aims to estimate the missing chromatic channels in grayscale images, generating diverse, colorful, and plausible results without relying on additional user-provided guidance. Chen et al. [8] pioneer the first deep-learning-based automatic colorization approach. With the advent of CNNs, which significantly increases the representation ability of the neural network, researchers have focused on developing adaptive feature extraction modules [10, 17, 23]. Recently, advanced generative models, e.g., VAEs [9], GANs [4, 15, 40, 41], cINNs [2], and Transformers [18, 22, 47], have been employed to better fit real-world color distributions and produce vivid colorization results. Concurrently, other works have explored incorporating various priors to achieve accurate colorization semantics, including segmentation masks [56, 55], detection boxes [39], rich color prior [44, 52], and pretrained generative models [20, 46]. Despite these advances, automatic colorization still faces challenges in handling multi-modal uncertainty. Considering that rich image semantics are encapsulated in language descriptions, we utilize the prior knowledge from the pretrained cross-modality generative model, i.e., Stable Diffusion [34], to alleviate this issue.
2.3 Text-driven image editing
Leveraging the flexibility of description representation, text-driven image editing methods enable precise control over image color, style, contour, and high-level semantics. Early works [24, 43, 51] develop text injection modules and design supervisory signals based on the conditional GAN [29]. Recently, with the development of diffusion models [13, 38], researchers show increasing interests in exploiting text-driven image generation [32, 33, 37], as well as diverse image editing tasks, e.g., text-guided image inpainting [34, 48], inversion-based style transfer [54], and text-driven image-to-image translation [19]. To reduce resource consumption, some researchers build models based on pretrained generative models by advanced sampling strategies [28, 30] or finetuning their weights [35, 50]. Although language-based colorization could be considered as a specialized text-driven image editing task, existing methods utilizing pretrained generative models hardly or seldom succeed in instance-aware colorization due to their failure to spatially align with input grayscale images.
In this section, we first provide a brief review of related diffusion models [13, 34, 38] in Sec. 3.1 to make this paper self-contained. Following that, we delve into the details of designed modules that spatially align with input conditions in Secs. 3.2 and 3.3. Next, we present the instance-aware sampling strategy for effective colorization in diverse and complex scenarios in Sec. 3.4. Finally, training details are shown in Sec. 3.5.
Diffusion models [13, 38] are one type of generative models that learn to represent the distribution of images from a given dataset. During the inference process, these models iteratively refine a
Figure 2: The framework of L-CAD. (a) We introduce an additional luminance encoder to capture multi-scale features from the grayscale image, which preserves local spatial structures of grayscale images and guides the decoding process from the latent space back to pixel space (Sec. 3.2). (b) We modify the noise prediction network by replacing Vanilla Convolution (VC) blocks with our proposed Channel-Extended Convolution (CEC) blocks within each downsampling module before Cross-Attention (CA) blocks. This prevents the ghosting effect by aligning with descriptions in the latent space (Sec. 3.3). For visualizing the colorized image, we further present the instance-aware sampling strategy to correctly assign colors to corresponding objects for descriptions with diverse levels of details and complex scenarios (Sec. 3.4). (c) During training, our model randomly receives any-level descriptions to increase the robustness of colorization (Sec. 3.5).
random noise sample  $x_{T} \sim \mathcal{N}(0,1)$  until it converges to a photorealistic image  $x_0$ . At each step  $t \in \{0,\dots,T\}$ , the intermediate sample  $\mathbf{x}_t$  is computed as:
$$
x _ {t} = \sqrt {\alpha_ {t}} x _ {0} + \sqrt {1 - \alpha_ {t}} \epsilon_ {t}, \tag {1}
$$
where  $\epsilon_t \sim \mathcal{N}(0,1)$  is the Gaussian noise at step  $t$ . In diffusion models, a neural network denoted as  $\epsilon_{\theta}$  is employed at each step to predict the noise  $\epsilon_t$ , guiding the model in the process of reversing  $x_{t}$  to  $x_{t-1}$  through a denoising procedure.
Diffusion models can be extended to handle conditional generative tasks by incorporating an additional condition  $y$ . In this case, the noise prediction network  $\epsilon_{\theta}$  takes three parameters as input: the noise sample  $x_{t}$ , the step  $t$ , and the condition  $y$ . This enables the diffusion process to generate results that are consistent with the given condition. To achieve the denoising objective, diffusion models typically utilize a simple  $\mathcal{L}_2$  loss, aiming to minimize the discrepancy between the learned image distribution and the target distribution. The loss function can be formulated as:
$$
\mathcal {L} _ {\mathrm {d m}} = \mathbb {E} _ {t, x _ {0}, \epsilon \sim \mathcal {N} (0, 1)} \left[ \| \epsilon_ {t} - \epsilon_ {\theta} (x _ {t}, t, y) \| ^ {2} \right]. \tag {2}
$$
To address the challenges of resource consumption and high-resolution image generation, Stable Diffusion [34] introduces a perceptual compression model to make the latent space perceptually equivalent to the image space. This enables the diffusion process to be performed in the compressed latent space. Specifically, it adopts a compression encoder  $\mathcal{E}$  to encode the given image  $x$  into latent space as  $z = \mathcal{E}(x)$ , and a compression decoder  $\mathcal{D}$  to reconstruct the image as  $\tilde{x} = \mathcal{D}(z)$ . Notably, the noise prediction network  $\epsilon_{\theta}$  is trained to learn the distribution of the latent representation  $z$  instead of the input image  $x$ .
In this paper, our goal is to leverage Stable Diffusion [34]'s robust language understanding and rich color priors for language-based colorization with any-level descriptions.
3.2 Luminance-guided image compression
While Stable Diffusion [34] demonstrates superior performance in text-driven image editing, it lacks the ability to preserve local spatial structures of input grayscale images, which is essential for the colorization task. To address this limitation, we propose the incorporation of an additional luminance
encoder in the pixel space. This encoder serves as a bridge between the colorization results and the grayscale images, ensuring better alignment between the two.
As shown in Fig. 2 (a), the colorful image is encoded into the latent representation by the compression encoder, while the luminance encoder extracts multi-scale features from grayscale images, preserving local structural semantics. These features are added directly into the corresponding scales of the compression decoder, guiding the decoding process from the latent space to the pixel space. The luminance encoder adopts the same architecture as the compression encoder, with the weights of the compression encoder and decoder fixed to retain prior knowledge from the pretrained model.
Inspired by LDL [25], we observe that erroneous pixels in regions with regular and sharp structures significantly damage visual perception. To address this, we intend to estimate an artifact map  $M_{h,w}^{\mathrm{art}}$ , which would serve to indicate the probability of encountering artifacts at specific spatial position  $(h,w)$  within the colorized results. Specifically, we calculate the residual  $\delta$  between the ground truth image  $x$  and the colorized result  $\tilde{x}$ . By computing the variance of the residual within local square windows at each position, we generate an artifact map denoted as  $M^{\mathrm{art}}$  as:
$$
M _ {h, w} ^ {\operatorname {a r t}} = \sum_ {p \in \Omega_ {(h, w)}} \left(\frac {\delta_ {p} - \mu_ {p}}{N _ {\mathrm {w i n}}}\right) ^ {2}, \quad \mu_ {p} = \sum_ {p \in \Omega_ {(h, w)}} \frac {\delta_ {p}}{N _ {\mathrm {w i n}} ^ {2}}, \tag {3}
$$
where  $p$  denotes the position index of the local window  $\Omega$  centered on position  $(h, w)$  and  $N_{\mathrm{win}}$  is the local window size. Given that artifacts typically show up with high-frequency characteristics, areas with higher variances likely indicate where these artifacts are. We could apply the artifact map as a weight to the image reconstruction loss as  $\mathcal{L}_{\mathrm{rec}} = \| M^{\mathrm{art}} \odot (x - \tilde{x}) \|_1$ .
Our total loss to train our model in pixel space is as follows:
$$
\mathcal {L} _ {\mathrm {p i x}} = \mathcal {L} _ {\mathrm {r e c}} + \alpha \mathcal {L} _ {\mathrm {d i s}} + \beta \mathcal {L} _ {\mathrm {p e r}}, \tag {4}
$$
where  $\alpha = 1.0$  and  $\beta = 0.5$  are weighting parameters, the discriminator loss  $\mathcal{L}_{\mathrm{dis}}$  and the perceptual loss  $\mathcal{L}_{\mathrm{per}}$  are adopted to improve the overall quality of reconstructed images.
3.3 Semantic-aligned latent representation
In Stable Diffusion [34], descriptions are incorporated into the latent representation using Cross-Attention (CA) blocks to ensure flexibility. However, this approach leads to the ghosting effect as color features are assigned outside the intended regions. To address this, we focus on maintaining semantic alignment between descriptions and grayscale images within the latent space.
In Fig. 2 (b), we replace Vanilla Convolution (VC) blocks in the noise prediction network with our Channel-Extended Convolution (CEC) blocks within each downsampling module. These CEC blocks receive the resized luminance features  $y^{\mathrm{lum}}$  from the luminance encoder in the pixel space as additional guidance. By leveraging extended channels, CEC blocks could effectively capture the local structural semantics of the luminance in the latent space. We formulate the channel-extended convolution as:
$$
f _ {h, w} ^ {\prime} = \sum_ {i = 0} ^ {N _ {\mathrm {k}} - 1} \sum_ {j = 0} ^ {N _ {\mathrm {k}} - 1} \left(\sum_ {k = 1} ^ {N _ {\mathrm {f i x}}} \omega_ {i, j, k} ^ {\mathrm {f i x}} f _ {p, q, k} + \sum_ {k = 1} ^ {N _ {\mathrm {e x t}}} \omega_ {i, j, N _ {\mathrm {f i x}} + k} ^ {\mathrm {e x t}} \bar {y} _ {p, q, k} ^ {\mathrm {l u m}}\right), \tag {5}
$$
where  $h, w, p = h + i - N_{\mathrm{k}} / 2$  and  $q = w + j - N_{\mathrm{k}} / 2$  are position index of feature maps,  $f$  and  $f'$  are separately feature maps before and after performing the convolution,  $N_{\mathrm{k}}$  is the kernel size,  $N_{\mathrm{fix}}$  and  $N_{\mathrm{ext}}$  separately mean the number of fixed channels and extended channels,  $\bar{y}^{\mathrm{lum}}$  denote resized luminance features. Mathematically, the CEC block is equivalent to using a stack of convolutions to extract features and add their output to the downsampling block. To preserve the robust language understanding and rich color priors of the pretrained generative model, we keep the weights  $\omega^{\mathrm{fix}}$  fixed during training. Additionally, we initialize the weights of extended channels  $\omega^{\mathrm{ext}}$  to zero, ensuring that our model maintains functional equivalence to the pretrained generative model prior to training.
We use skip connections between downsampling and upsampling modules to guide the upsampling process with luminance features. To expedite training convergence, we utilize Channel-Extended Convolution (CEC) blocks only in downsampling modules.
Table 1: Quantitative comparisons with language-based colorization methods (left) and automatic colorization methods (right). * marks our method receiving scarce-level descriptions. Throughout this paper,  $\uparrow \left(  \downarrow  \right)$  means higher (lower) is better. Best performances are highlighted in bold.
Comparison with language-based colorization methods	Comparison with automatic colorization methods
Method	Extended COCO-Stuff	Multi-instance	Method	Extended COCO-Stuff	Multi-instance
PSNR↑	SSIM↑	LPIPS↓	PSNR↑	SSIM↑	LPIPS↓	PSNR↑	SSIM↑	LPIPS↓	PSNR↑	SSIM↑	LPIPS↓
LBIE	22.02	.8519	.265	21.92	.8697	.260	CIC	22.03	.8893	.232	21.95	.8754	.239
ML2018	21.06	.8533	.282	20.54	.8495	.294	InstColor	23.76	.8979	.195	23.59	.8968	.199
Xie2018	21.92	.8764	.233	20.04	.8530	.268	ChromaGAN	22.08	.8416	.275	22.41	.8584	.248
L-CoDe	24.86	.9091	.166	23.84	.9056	.176	BigColor	21.45	.8868	.236	21.43	.8824	.225
L-CoDer	25.03	.9125	.162	24.16	.9135	.168	DISCO	20.62	.8744	.212	20.81	.8762	.204
L-CoIns	25.10	.9151	.159	24.74	.9141	.156	CT²	23.16	.9030	.191	22.65	.8993	.193
W/o LIC	23.87	.8866	.190	23.96	.9047	.160	W/o LIC*	22.25	.8584	.206	22.53	.8697	.176
W/o SLR	17.26	.5984	.587	15.61	.4879	.531	W/o SLR*	15.73	.6132	.512	14.18	.5732	.472
W/o ISS	25.32	.9174	.147	24.57	.9138	.149	W/o ISS*	24.24	.9073	.159	24.07	.9054	.165
L-CAD	25.97	.9191	.142	25.51	.9156	.127	L-CAD*	24.33	.9114	.157	24.19	.9108	.164
The loss for training the denoising network in the latent space is:
$$
\mathcal {L} _ {\text {l a t}} = \mathbb {E} _ {t, z _ {0}, \epsilon \sim \mathcal {N} (0, 1)} \left[ \| \epsilon_ {t} - \epsilon_ {\theta} \left(z _ {t}, t, y ^ {\text {t e x}}, y ^ {\text {l u m}}\right) \| ^ {2} \right], \tag {6}
$$
where  $y^{\mathrm{tex}}$  is any-level descriptions encoded by the text encoder of CLIP [31].
3.4 Instance-aware sampling strategy
To ensure accurate color assignment to objects in images with diverse and complex descriptions, we propose an instance-aware sampling strategy. This strategy draws inspiration from the effectiveness of InstColor [39], which employs an auxiliary recognition model to enhance instance-aware performance. By incorporating the instance-aware sampling strategy into our colorization process, we aim to improve the accuracy and effectiveness of assigning colors to corresponding objects.
To improve instance-awareness, we employ a referring segmentation model (e.g., SAM [21]) to estimate object contours mentioned in the description  $(M^{\mathrm{est}})$ . Attention maps  $(M_{l}^{\mathrm{att}})$  at the  $l$ -th CA block control colorized regions for each color word. We normalize the attention maps with Sigmoid and refine them iteratively by aligning with downsampled estimated contours. After that, we apply the denoising process in DDIM [38]. Algorithm 1 outlines our instance-aware sampling strategy.
Algorithm 1: Instance-aware sampling strategy
input: Roughly estimated object contour  $M^{\mathrm{est}}$   
output: Colorized latent representation  $z_0$   
for  $t = T\ldots 1$  do  
 $\begin{array}{rl} & M_{*}^{\mathrm{att}} = \epsilon_{\theta}(z_t,t,y^{\mathrm{lum}},y^{\mathrm{tex}})\\ & \mathbf{for} l = 1\dots L\mathbf{do}\\ & \hat{M}_l^{\mathrm{est}}\gets \mathrm{Downsampling}(M^{\mathrm{est}},l)\\ & \mathcal{M}\gets \mathrm{Sigmoid}(M_l^{\mathrm{att}})\\ & \hat{M}_l^{\mathrm{att}}\gets M_l^{\mathrm{att}} - \lambda \nabla_{\mathcal{M}}\mathcal{L}_{\mathrm{BCE}}(\mathcal{M},\hat{M}_l^{\mathrm{est}}) \end{array}$   
end  
 $\hat{\epsilon}_{t, - } = \epsilon_{\theta}(z_{t},t,y^{\mathrm{lum}},y^{\mathrm{tex}})\{M_{*}^{\mathrm{att}}\leftarrow \hat{M}_{*}^{\mathrm{att}}\}$ $z_{t - 1} = \mathrm{DDIM}(z_t,\hat{\epsilon}_t,t)$
Our instance-aware sampling strategy does not heavily rely on precise segmentation results from the referring segmentation model. This is due to the following reasons: (i) Our sampling is performed in the latent space using low-resolution segmentation results, mitigating the impact of segmentation imprecision. (ii) When decoding the latent representation into the pixel space, the compression decoder, guided by local structural semantics from grayscale images, could adaptively fix color bleeding issues to produce visually pleasing results.
3.5 Learning and implementation details
During training, we utilize the classifier-free guidance [14] to supervise the colorization process conditioned on descriptions. To increase robustness for any-level descriptions, we randomly replace  $30\%$  of the complete-level and partial-level descriptions with scarce-level descriptions. The pixel space modules are trained using Eq. (4) with a batch size of 8 for 10 epochs until convergence.
Figure 3: Compared with language-based colorization methods, our model achieves correct instance-aware colorization while improving the colorized quality for unmentioned objects.
Figure 4: Compared with automatic colorization methods, our model shows comparable performance with scarce-level descriptions that lack meaningful color information.