add implementation scripts

app.py

@@ -0,0 +1,135 @@
+"""
+Small demo application to explore Gradio.
+"""
+
+import argparse
+import os
+from functools import partial
+
+import gradio as gr
+from PIL import Image
+
+from die_model import UNetDIEModel
+from utils import resize_image, make_image_square, cast_pil_image_to_torch_tensor_with_4_channel_dim, \
+    remove_square_padding
+
+
+def die_inference(
+    image_raw,
+    num_of_die_iterations,
+    die_model,
+    device
+):
+    """
+    Function to run the DIE model.
+    :param image_raw: raw image
+    :param num_of_die_iterations: number of DIE iterations
+    :param die_model: DIE model
+    :param device: device
+    :return: cleaned image
+    """
+
+    # preprocess
+    image_raw_resized = resize_image(image_raw, 1500)
+    image_raw_resized_square = make_image_square(image_raw_resized)
+    image_raw_resized_square_tensor = cast_pil_image_to_torch_tensor_with_4_channel_dim(image_raw_resized_square)
+    image_raw_resized_square_tensor = image_raw_resized_square_tensor.to(device)
+
+    # convert string to int
+    num_of_die_iterations = int(num_of_die_iterations)
+
+    # inference
+    image_die = die_model.enhance_document_image(
+        image_raw_list=[image_raw_resized_square_tensor],
+        num_of_die_iterations=num_of_die_iterations
+    )[0]
+
+    # postprocess
+    image_die_resized = remove_square_padding(
+        original_image=image_raw,
+        square_image=image_die,
+        resize_back_to_original=True
+    )
+
+
+    return image_die_resized
+
+
+def main():
+    """
+    Main function to run the Gradio demo.
+    :return:
+    """
+
+    args = parse_arguments()
+
+    description = "Welcome to the Document Image Enhancement (DIE) model demo on Hugging Face!\n\n" \
+                  "" \
+                  "This interactive application showcases a specialized AI model developed by " \
+                  "the [Artificial Intelligence group](https://ai.renyi.hu) at the [Alfréd Rényi Institute of Mathematics](https://renyi.hu).\n\n" \
+                  "" \
+                  "Our DIE model is designed to enhance and restore archival and aged document images " \
+                  "by removing various types of degradation, thereby making historical documents more legible " \
+                  "and suitable for Optical Character Recognition (OCR) processing.\n\n" \
+                  "" \
+                  "The model effectively tackles 20-30 types of domain-specific noise found in historical records, " \
+                  "such as scribbles, bleed-through text, faded or worn text, blurriness, textured noise, " \
+                  "and unwanted background elements. " \
+                  "By applying deep learning techniques, specifically a U-Net-based architecture, " \
+                  "the model accurately cleans and clarifies text while preserving original details. " \
+                  "This improved clarity dramatically boosts OCR accuracy, making it an ideal " \
+                  "pre-processing tool in digitization workflows.\n\n" \
+                  "" \
+                  "If you’re interested in learning more about the model’s capabilities or potential applications, " \
+                  "please contact us at: [email protected].\n\n"
+
+    # TODO: Add a description for the Number of DIE iterations parameter!
+
+    num_of_die_iterations_list = [1, 2, 3]
+
+    # Provide images alone for example display
+    example_image_list = [
+        [Image.open(os.path.join(args.example_image_path, image_path))]
+        for image_path in os.listdir(args.example_image_path)
+    ]
+
+    # Load DIE model
+    die_model = UNetDIEModel(args=args)
+
+    # Partially apply the model and device arguments to die_inference
+    partial_die_inference = partial(die_inference, device=args.device, die_model=die_model)
+
+    demo = gr.Interface(
+        fn=partial_die_inference,
+        inputs=[
+            gr.Image(type="pil", label="Degraded Document Image"),
+            gr.Dropdown(num_of_die_iterations_list, label="Number of DIE iterations", value=1),
+        ],
+        outputs=gr.Image(type="pil", label="Clean Document Image"),
+        title="Document Image Enhancement (DIE) model",
+        description=description,
+        examples=example_image_list
+    )
+
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+
+
+def parse_arguments():
+    """
+    Parse arguments.
+    :return: argument namespace
+    """
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--die_model_path", default="./2024_08_09_model_epoch_89.pt")
+    parser.add_argument("--device", default="cpu")
+
+    parser.add_argument("--example_image_path", default="./example_images")
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+
+    main()

die_model.py

@@ -0,0 +1,264 @@
+"""
+U-Net based DIE model for cleaning document.
+"""
+
+import os
+from typing import Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+from PIL import Image
+
+
+class DoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        # if you have padding issues, see
+        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = torch.sigmoid(x)
+        return x
+
+
+class UNet(nn.Module):
+    def __init__(self, n_channels, output_channel_dim=1, bilinear=False):
+        super(UNet, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = output_channel_dim
+        self.bilinear = bilinear
+
+        self.inc = DoubleConv(n_channels, 64)
+        self.down1 = Down(64, 128)
+        self.down2 = Down(128, 256)
+        self.down3 = Down(256, 512)
+        factor = 2 if bilinear else 1
+        self.down4 = Down(512, 1024 // factor)
+        self.up1 = Up(1024, 512 // factor, bilinear)
+        self.up2 = Up(512, 256 // factor, bilinear)
+        self.up3 = Up(256, 128 // factor, bilinear)
+        self.up4 = Up(128, 64, bilinear)
+        self.outc = OutConv(64, output_channel_dim)
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        logits = self.outc(x)
+        return logits
+
+
+def add_gaussian_noise(
+    data: torch.Tensor
+) -> torch.Tensor:
+    """
+    Adding gaussian noise to torch tensor.
+    :param data: torch tensor
+    :return: noise perturbed tensor
+    """
+
+    data_with_noise = data.clone()
+    data_with_noise += torch.normal(mean=0, std=0.05, size=data_with_noise.shape).to(data_with_noise.device)
+    data_with_noise = data_with_noise.clip(min=0, max=1)
+
+    return data_with_noise
+
+
+def inference_model(
+    model: Callable,
+    model_input: torch.Tensor,
+    device: str | torch.device,
+    num_of_iterations: int = 1
+) -> list[torch.Tensor, ...]:
+    """
+    Performing model inference.
+    :param model: image pre-processing model
+    :param model_input: data to model
+    :param device: cuda device
+    :param num_of_iterations: defines how many times feed the network (recursively)
+    :return: predictions
+    """
+
+    # inference model
+    with torch.no_grad():
+
+        prediction_list = []
+
+        model_input = model_input.to(device)
+
+        if len(model_input.shape) == 3:
+            model_input = model_input.unsqueeze(dim=0)
+
+        model_input_original_part = model_input[:, 0:3, ...]
+
+        for i in range(num_of_iterations):
+
+            if i == 0:
+                model_input = add_gaussian_noise(model_input)
+                prediction = model(model_input)
+                prediction_list.append(prediction)
+                model_input_new = torch.cat((model_input_original_part, prediction.detach()), dim=1)
+            else:
+                model_input_perturbed = add_gaussian_noise(model_input_new)
+                prediction = model(model_input_perturbed)
+                prediction_list.append(prediction)
+                model_input_new = torch.cat((model_input_original_part, prediction.detach()), dim=1)
+
+    return prediction_list
+
+
+def load_unet(
+    model_path: str,
+    device: str = 'cpu',
+    eval_mode: bool = False,
+    n_channels: int = 4,
+    bilinear: bool = False,
+    output_channel_dim: int = 1
+):
+
+    print("Loading UNet model...")
+
+    # image preprocessing model
+    model = UNet(
+        n_channels=n_channels,
+        bilinear=bilinear,
+        output_channel_dim=output_channel_dim
+    )
+
+    # this hack is required due to distributed data parallel training
+    state_dict = torch.load(os.path.join(model_path), map_location=device)
+    new_state_dict = {key.replace('module.', ''): value for key, value in state_dict.items()}
+    model.load_state_dict(new_state_dict)
+    model.to(device)
+
+    if eval_mode:
+        model.eval()
+
+    return model
+
+
+class UNetDIEModel:
+    """
+    Class for Document Image Enhancement with U-Net.
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs
+    ):
+        """
+        Initialization.
+        """
+
+        self.args = kwargs['args']
+
+        # loading text detector model
+        self.die = load_unet(
+            model_path=self.args.die_model_path,
+            device=self.args.device,
+            eval_mode=True,
+        )
+
+    def enhance_document_image(
+        self,
+        image_raw_list: list[Image.Image],
+        num_of_die_iterations: int = 1,
+    ) -> list[Image.Image]:
+        """"
+        Enhance document image by removing noise.
+        :param image_raw_list: original document page to process
+        :param num_of_die_iterations: number of DIE iterations
+        :return: cleaned document page to process
+        """
+
+        with torch.no_grad():
+
+            # image_die = torch.stack(image_die_list, dim=0)
+            image_die = torch.stack(image_raw_list, dim=0)
+
+            # document image enhancement
+            prediction_list = inference_model(
+                model=self.die,
+                model_input=image_die,
+                num_of_iterations=num_of_die_iterations,
+                device=self.args.device
+            )
+
+            # transform DIE model output to image and apply post-processing
+            last_prediction = prediction_list[-1]
+            batch_size = last_prediction.size(0)
+            image_die_list = [T.ToPILImage()(last_prediction[idx, ...]).convert('RGB') for idx in range(batch_size)]
+
+            return image_die_list

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+pillow
+torch
+torchvision

utils.py

@@ -0,0 +1,139 @@
+"""
+Utility functions for the DIE demo.
+"""
+
+
+import torch
+from PIL import Image
+from torch import Tensor
+from torchvision import transforms
+
+
+def resize_image(
+    image: Image.Image,
+    max_size: int = 1024
+) -> Image.Image:
+    """
+    Resizing images by keeping the ratios
+    :param image: PIL image
+    :param max_size: size of the new image larger side
+    :return: the resized PIL image
+    """
+
+    # extracting size
+    width, height = image.size
+
+    # checking which side is larger
+    height_larger = True if height >= width else False
+
+    # reshaping based on the larger side
+    if height_larger:
+        height_new = max_size
+        width_new = round((height_new / height) * width)
+    else:
+        width_new = max_size
+        height_new = round((width_new / width) * height)
+
+    return image.resize((width_new, height_new))
+
+
+def make_image_square(
+    image: Image.Image,
+    image_size: int = 1024
+) -> Image.Image:
+    """
+    Making the input image a square
+    :param image: PIL image
+    :param image_size: defines the size of the square image
+    :return: the square-sized PIL image
+    """
+
+    if max(image.size) > image_size:
+        image_size = max(image.size)
+    # creating a new square image
+    if image.mode == 'L':
+        image_square = Image.new(image.mode, (image_size, image_size), (255,))
+    elif image.mode == 'RGB':
+        image_square = Image.new(image.mode, (image_size, image_size), (255, 255, 255))
+    else:
+        raise NotImplementedError("Not implemented image mode.")
+    # copying the original content onto the blank image
+    image_square.paste(image, (0, 0))
+
+    return image_square
+
+
+def cast_pil_image_to_torch_tensor_with_4_channel_dim(
+    image: Image.Image,
+    device: str | None = None
+) -> Tensor:
+    """
+    Casting PIL image to torch tensor.
+    Adding the grayscale image of the original RGB image as a 4th channel dimension.
+    :param image: input image
+    :param device: cuda device
+    :return: torch tensor (4 channel dim)
+    """
+
+    # PIL image to torch tensor transformation
+    transform = transforms.Compose([transforms.PILToTensor()])
+
+    # creating gray image
+    image_gray = image.convert('L')
+
+    # casting PIL images to torch tensor with normalization
+    image_tensor = transform(image.convert('RGB')).to(torch.float32) / 255.0
+    image_gray_tensor = transform(image_gray).to(torch.float32) / 255.0
+
+    # concatenating gray channel to RGB channel
+    final_image_tensor = torch.cat((image_tensor, image_gray_tensor), dim=0)
+
+    # moving tensor to gpu if required
+    if device is not None:
+        final_image_tensor = final_image_tensor.to(device)
+
+    return final_image_tensor
+
+
+def remove_square_padding(
+    original_image: Image.Image | Tensor,
+    square_image: Image.Image | Tensor,
+    resize_back_to_original: bool = False
+):
+    """
+    Removing the square padding added to the original image to make square.
+    :param original_image: the image with the original size
+    :param square_image: the image with the square size
+    :param resize_back_to_original: defines if we want to resize the square image back to the original size
+    :return: square image with the original size ratio
+    """
+
+    if isinstance(original_image, Image.Image):
+        original_width, original_height = original_image.size
+    else:
+        original_height, original_width = original_image.shape[:2]
+
+    if isinstance(square_image, Image.Image):
+        square_width, square_height = square_image.size
+    else:
+        square_height, square_width = square_image.shape[:2]
+
+    if original_width > original_height:
+        ratio = square_width / original_width
+        new_width = square_width
+        new_height = int(ratio * original_height)
+    else:
+        ratio = square_height / original_height
+        new_height = square_height
+        new_width = int(ratio * original_width)
+
+    # cutting size of the square image to the original ratio
+    if isinstance(square_image, Image.Image):
+        square_image_with_original_ratio = square_image.crop((0, 0, new_width, new_height))
+    else:
+        square_image_with_original_ratio = square_image[:new_height, :new_width]
+
+    if resize_back_to_original:
+        square_image_with_original_ratio = square_image_with_original_ratio.resize((original_width, original_height))
+
+    return square_image_with_original_ratio