initial commit

.gitignore

@@ -0,0 +1,252 @@
+# Created by https://www.toptal.com/developers/gitignore/api/macos,windows,linux,python
+# Edit at https://www.toptal.com/developers/gitignore?templates=macos,windows,linux,python
+
+### Linux ###
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# KDE directory preferences
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+### macOS ###
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+### macOS Patch ###
+# iCloud generated files
+*.icloud
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+### Windows ###
+# Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+
+# End of https://www.toptal.com/developers/gitignore/api/macos,windows,linux,python
+
+*.pth
+gradio_cached_examples

README.md

@@ -1,6 +1,6 @@
 ---
-title: MangaLineExtraction
-emoji: 💻
+title: Anime to Sketch
+emoji: 💭
 colorFrom: green
 colorTo: yellow
 sdk: gradio
@@ -10,4 +10,7 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Original repo:
+- MangaLineExtraction: https://github.com/ljsabc/MangaLineExtraction_PyTorch
+- Anime2Sketch: https://github.com/Mukosame/Anime2Sketch
+

anime2sketch/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Xiaoyu Xiang
+
+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.

anime2sketch/model.py

@@ -0,0 +1,256 @@
+import os
+from PIL import Image
+import torchvision.transforms as transforms
+
+try:
+    from torchvision.transforms import InterpolationMode
+
+    bic = InterpolationMode.BICUBIC
+except ImportError:
+    bic = Image.BICUBIC
+
+import numpy as np
+import torch
+import torch.nn as nn
+import functools
+
+IMG_EXTENSIONS = [".jpg", ".jpeg", ".png", ".ppm", ".bmp", ".webp"]
+
+
+class UnetGenerator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(
+        self,
+        input_nc,
+        output_nc,
+        num_downs,
+        ngf=64,
+        norm_layer=nn.BatchNorm2d,
+        use_dropout=False,
+    ):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(UnetGenerator, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 8,
+            ngf * 8,
+            input_nc=None,
+            submodule=None,
+            norm_layer=norm_layer,
+            innermost=True,
+        )  # add the innermost layer
+        for _ in range(num_downs - 5):  # add intermediate layers with ngf * 8 filters
+            unet_block = UnetSkipConnectionBlock(
+                ngf * 8,
+                ngf * 8,
+                input_nc=None,
+                submodule=unet_block,
+                norm_layer=norm_layer,
+                use_dropout=use_dropout,
+            )
+        # gradually reduce the number of filters from ngf * 8 to ngf
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+        )
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+        )
+        unet_block = UnetSkipConnectionBlock(
+            ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+        )
+        self.model = UnetSkipConnectionBlock(
+            output_nc,
+            ngf,
+            input_nc=input_nc,
+            submodule=unet_block,
+            outermost=True,
+            norm_layer=norm_layer,
+        )  # add the outermost layer
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class UnetSkipConnectionBlock(nn.Module):
+    """Defines the Unet submodule with skip connection.
+    X -------------------identity----------------------
+    |-- downsampling -- |submodule| -- upsampling --|
+    """
+
+    def __init__(
+        self,
+        outer_nc,
+        inner_nc,
+        input_nc=None,
+        submodule=None,
+        outermost=False,
+        innermost=False,
+        norm_layer=nn.BatchNorm2d,
+        use_dropout=False,
+    ):
+        """Construct a Unet submodule with skip connections.
+        Parameters:
+            outer_nc (int) -- the number of filters in the outer conv layer
+            inner_nc (int) -- the number of filters in the inner conv layer
+            input_nc (int) -- the number of channels in input images/features
+            submodule (UnetSkipConnectionBlock) -- previously defined submodules
+            outermost (bool)    -- if this module is the outermost module
+            innermost (bool)    -- if this module is the innermost module
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+        """
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(
+            input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias
+        )
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(
+                inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1
+            )
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(
+                inner_nc, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias
+            )
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(
+                inner_nc * 2,
+                outer_nc,
+                kernel_size=4,
+                stride=2,
+                padding=1,
+                bias=use_bias,
+            )
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:  # add skip connections
+            return torch.cat([x, self.model(x)], 1)
+
+
+class Anime2Sketch:
+    def __init__(
+        self, model_path: str = "./models/netG.pth", device: str = "cpu"
+    ) -> None:
+        norm_layer = functools.partial(
+            nn.InstanceNorm2d, affine=False, track_running_stats=False
+        )
+        net = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False)
+        ckpt = torch.load(model_path)
+
+        for key in list(ckpt.keys()):
+            if "module." in key:
+                ckpt[key.replace("module.", "")] = ckpt[key]
+                del ckpt[key]
+
+        net.load_state_dict(ckpt)
+
+        self.model = net
+
+        if torch.cuda.is_available() and device == "cuda":
+            self.device = "cuda"
+            self.model.to(device)
+        else:
+            self.device = "cpu"
+            self.model.to("cpu")
+
+    def predict(self, image: Image.Image, load_size: int = 512) -> Image:
+        try:
+            aus_resize = None
+            if load_size > 0:
+                aus_resize = image.size
+            transform = self.get_transform(load_size=load_size)
+            image = transform(image)
+            img = image.unsqueeze(0)
+        except:
+            raise Exception("Error in reading image {}".format(image.filename))
+
+        aus_tensor = self.model(img.to(self.device))
+        aus_img = self.tensor_to_img(aus_tensor)
+
+        image_pil = Image.fromarray(aus_img)
+        if aus_resize:
+            bic = Image.BICUBIC
+            image_pil = image_pil.resize(aus_resize, bic)
+
+        return image_pil
+
+    def get_transform(self, load_size=0, grayscale=False, method=bic, convert=True):
+        transform_list = []
+        if grayscale:
+            transform_list.append(transforms.Grayscale(1))
+        if load_size > 0:
+            osize = [load_size, load_size]
+            transform_list.append(transforms.Resize(osize, method))
+        if convert:
+            transform_list += [transforms.ToTensor()]
+            if grayscale:
+                transform_list += [transforms.Normalize((0.5,), (0.5,))]
+            else:
+                transform_list += [
+                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+                ]
+        return transforms.Compose(transform_list)
+
+    def tensor_to_img(self, input_image, imtype=np.uint8):
+        """ "Converts a Tensor array into a numpy image array.
+        Parameters:
+            input_image (tensor) --  the input image tensor array
+            imtype (type)        --  the desired type of the converted numpy array
+        """
+
+        if not isinstance(input_image, np.ndarray):
+            if isinstance(input_image, torch.Tensor):  # get the data from a variable
+                image_tensor = input_image.data
+            else:
+                return input_image
+            image_numpy = (
+                image_tensor[0].cpu().float().numpy()
+            )  # convert it into a numpy array
+            if image_numpy.shape[0] == 1:  # grayscale to RGB
+                image_numpy = np.tile(image_numpy, (3, 1, 1))
+            image_numpy = (
+                (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
+            )  # post-processing: tranpose and scaling
+        else:  # if it is a numpy array, do nothing
+            image_numpy = input_image
+        return image_numpy.astype(imtype)

app.py

@@ -0,0 +1,85 @@
+import gradio as gr
+from setup import setup
+import cv2
+from PIL import Image
+from manga_line_extraction.model import MangaLineExtractor
+from anime2sketch.model import Anime2Sketch
+
+setup()
+
+print("Setup finished")
+
+extractor = MangaLineExtractor("./models/erika.pth", "cpu")
+to_sketch = Anime2Sketch("./models/netG.pth", "cpu")
+
+print("Model loaded")
+
+
+def extract(image):
+    return extractor.predict(image)
+
+
+def convert_to_sketch(image):
+    return to_sketch.predict(image)
+
+
+def start(image):
+    return [extract(image), convert_to_sketch(Image.fromarray(image).convert("RGB"))]
+
+
+def ui():
+    with gr.Blocks() as blocks:
+        gr.Markdown(
+            """
+        # Anime to Sketch 
+        Unofficial demo for converting illustrations into sketches. 
+        Original repos:
+        - [MangaLineExtraction_PyTorch](https://github.com/ljsabc/MangaLineExtraction_PyTorch)
+        - [Anime2Sketch](https://github.com/Mukosame/Anime2Sketch)
+        """
+        )
+
+        with gr.Row():
+            with gr.Column():
+                input_img = gr.Image(label="Input", interactive=True)
+
+                extract_btn = gr.Button("Extract", variant="primary")
+
+            with gr.Column():
+                # with gr.Row():
+                extract_output_img = gr.Image(
+                    label="MangaLineExtraction", interactive=False
+                )
+                to_sketch_output_img = gr.Image(label="Anime2Sketch", interactive=False)
+
+        gr.Examples(
+            fn=start,
+            examples=[
+                ["./examples/1.jpg"],
+                ["./examples/2.jpg"],
+                ["./examples/3.jpg"],
+                ["./examples/4.jpg"],
+                ["./examples/5.jpg"],
+                ["./examples/6.jpg"],
+                ["./examples/7.jpg"],
+                ["./examples/8.jpg"],
+            ],
+            inputs=[input_img],
+            outputs=[extract_output_img, to_sketch_output_img],
+            label="Examples",
+            cache_examples=True,
+        )
+
+        gr.Markdown("Images are from nijijourney.")
+
+        extract_btn.click(
+            fn=start,
+            inputs=[input_img],
+            outputs=[extract_output_img, to_sketch_output_img],
+        )
+
+    return blocks
+
+
+if __name__ == "__main__":
+    ui().launch()

manga_line_extraction/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Miaomiao Li
+
+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.

manga_line_extraction/model.py

@@ -0,0 +1,323 @@
+import os
+import torch
+import torch.nn as nn
+from torch.utils.data.dataset import Dataset
+from PIL import Image
+import fnmatch
+import cv2
+
+import sys
+
+import numpy as np
+
+# torch.set_printoptions(precision=10)
+
+
+class _bn_relu_conv(nn.Module):
+    def __init__(self, in_filters, nb_filters, fw, fh, subsample=1):
+        super(_bn_relu_conv, self).__init__()
+        self.model = nn.Sequential(
+            nn.BatchNorm2d(in_filters, eps=1e-3),
+            nn.LeakyReLU(0.2),
+            nn.Conv2d(
+                in_filters,
+                nb_filters,
+                (fw, fh),
+                stride=subsample,
+                padding=(fw // 2, fh // 2),
+                padding_mode="zeros",
+            ),
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+        # the following are for debugs
+        print(
+            "****",
+            np.max(x.cpu().numpy()),
+            np.min(x.cpu().numpy()),
+            np.mean(x.cpu().numpy()),
+            np.std(x.cpu().numpy()),
+            x.shape,
+        )
+        for i, layer in enumerate(self.model):
+            if i != 2:
+                x = layer(x)
+            else:
+                x = layer(x)
+                # x = nn.functional.pad(x, (1, 1, 1, 1), mode='constant', value=0)
+            print(
+                "____",
+                np.max(x.cpu().numpy()),
+                np.min(x.cpu().numpy()),
+                np.mean(x.cpu().numpy()),
+                np.std(x.cpu().numpy()),
+                x.shape,
+            )
+            print(x[0])
+        return x
+
+
+class _u_bn_relu_conv(nn.Module):
+    def __init__(self, in_filters, nb_filters, fw, fh, subsample=1):
+        super(_u_bn_relu_conv, self).__init__()
+        self.model = nn.Sequential(
+            nn.BatchNorm2d(in_filters, eps=1e-3),
+            nn.LeakyReLU(0.2),
+            nn.Conv2d(
+                in_filters,
+                nb_filters,
+                (fw, fh),
+                stride=subsample,
+                padding=(fw // 2, fh // 2),
+            ),
+            nn.Upsample(scale_factor=2, mode="nearest"),
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class _shortcut(nn.Module):
+    def __init__(self, in_filters, nb_filters, subsample=1):
+        super(_shortcut, self).__init__()
+        self.process = False
+        self.model = None
+        if in_filters != nb_filters or subsample != 1:
+            self.process = True
+            self.model = nn.Sequential(
+                nn.Conv2d(in_filters, nb_filters, (1, 1), stride=subsample)
+            )
+
+    def forward(self, x, y):
+        # print(x.size(), y.size(), self.process)
+        if self.process:
+            y0 = self.model(x)
+            # print("merge+", torch.max(y0+y), torch.min(y0+y),torch.mean(y0+y), torch.std(y0+y), y0.shape)
+            return y0 + y
+        else:
+            # print("merge", torch.max(x+y), torch.min(x+y),torch.mean(x+y), torch.std(x+y), y.shape)
+            return x + y
+
+
+class _u_shortcut(nn.Module):
+    def __init__(self, in_filters, nb_filters, subsample):
+        super(_u_shortcut, self).__init__()
+        self.process = False
+        self.model = None
+        if in_filters != nb_filters:
+            self.process = True
+            self.model = nn.Sequential(
+                nn.Conv2d(
+                    in_filters,
+                    nb_filters,
+                    (1, 1),
+                    stride=subsample,
+                    padding_mode="zeros",
+                ),
+                nn.Upsample(scale_factor=2, mode="nearest"),
+            )
+
+    def forward(self, x, y):
+        if self.process:
+            return self.model(x) + y
+        else:
+            return x + y
+
+
+class basic_block(nn.Module):
+    def __init__(self, in_filters, nb_filters, init_subsample=1):
+        super(basic_block, self).__init__()
+        self.conv1 = _bn_relu_conv(
+            in_filters, nb_filters, 3, 3, subsample=init_subsample
+        )
+        self.residual = _bn_relu_conv(nb_filters, nb_filters, 3, 3)
+        self.shortcut = _shortcut(in_filters, nb_filters, subsample=init_subsample)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.residual(x1)
+        return self.shortcut(x, x2)
+
+
+class _u_basic_block(nn.Module):
+    def __init__(self, in_filters, nb_filters, init_subsample=1):
+        super(_u_basic_block, self).__init__()
+        self.conv1 = _u_bn_relu_conv(
+            in_filters, nb_filters, 3, 3, subsample=init_subsample
+        )
+        self.residual = _bn_relu_conv(nb_filters, nb_filters, 3, 3)
+        self.shortcut = _u_shortcut(in_filters, nb_filters, subsample=init_subsample)
+
+    def forward(self, x):
+        y = self.residual(self.conv1(x))
+        return self.shortcut(x, y)
+
+
+class _residual_block(nn.Module):
+    def __init__(self, in_filters, nb_filters, repetitions, is_first_layer=False):
+        super(_residual_block, self).__init__()
+        layers = []
+        for i in range(repetitions):
+            init_subsample = 1
+            if i == repetitions - 1 and not is_first_layer:
+                init_subsample = 2
+            if i == 0:
+                l = basic_block(
+                    in_filters=in_filters,
+                    nb_filters=nb_filters,
+                    init_subsample=init_subsample,
+                )
+            else:
+                l = basic_block(
+                    in_filters=nb_filters,
+                    nb_filters=nb_filters,
+                    init_subsample=init_subsample,
+                )
+            layers.append(l)
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class _upsampling_residual_block(nn.Module):
+    def __init__(self, in_filters, nb_filters, repetitions):
+        super(_upsampling_residual_block, self).__init__()
+        layers = []
+        for i in range(repetitions):
+            l = None
+            if i == 0:
+                l = _u_basic_block(
+                    in_filters=in_filters, nb_filters=nb_filters
+                )  # (input)
+            else:
+                l = basic_block(in_filters=nb_filters, nb_filters=nb_filters)  # (input)
+            layers.append(l)
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class res_skip(nn.Module):
+    def __init__(self):
+        super(res_skip, self).__init__()
+        self.block0 = _residual_block(
+            in_filters=1, nb_filters=24, repetitions=2, is_first_layer=True
+        )  # (input)
+        self.block1 = _residual_block(
+            in_filters=24, nb_filters=48, repetitions=3
+        )  # (block0)
+        self.block2 = _residual_block(
+            in_filters=48, nb_filters=96, repetitions=5
+        )  # (block1)
+        self.block3 = _residual_block(
+            in_filters=96, nb_filters=192, repetitions=7
+        )  # (block2)
+        self.block4 = _residual_block(
+            in_filters=192, nb_filters=384, repetitions=12
+        )  # (block3)
+
+        self.block5 = _upsampling_residual_block(
+            in_filters=384, nb_filters=192, repetitions=7
+        )  # (block4)
+        self.res1 = _shortcut(
+            in_filters=192, nb_filters=192
+        )  # (block3, block5, subsample=(1,1))
+
+        self.block6 = _upsampling_residual_block(
+            in_filters=192, nb_filters=96, repetitions=5
+        )  # (res1)
+        self.res2 = _shortcut(
+            in_filters=96, nb_filters=96
+        )  # (block2, block6, subsample=(1,1))
+
+        self.block7 = _upsampling_residual_block(
+            in_filters=96, nb_filters=48, repetitions=3
+        )  # (res2)
+        self.res3 = _shortcut(
+            in_filters=48, nb_filters=48
+        )  # (block1, block7, subsample=(1,1))
+
+        self.block8 = _upsampling_residual_block(
+            in_filters=48, nb_filters=24, repetitions=2
+        )  # (res3)
+        self.res4 = _shortcut(
+            in_filters=24, nb_filters=24
+        )  # (block0,block8, subsample=(1,1))
+
+        self.block9 = _residual_block(
+            in_filters=24, nb_filters=16, repetitions=2, is_first_layer=True
+        )  # (res4)
+        self.conv15 = _bn_relu_conv(
+            in_filters=16, nb_filters=1, fh=1, fw=1, subsample=1
+        )  # (block7)
+
+    def forward(self, x):
+        x0 = self.block0(x)
+        x1 = self.block1(x0)
+        x2 = self.block2(x1)
+        x3 = self.block3(x2)
+        x4 = self.block4(x3)
+
+        x5 = self.block5(x4)
+        res1 = self.res1(x3, x5)
+
+        x6 = self.block6(res1)
+        res2 = self.res2(x2, x6)
+
+        x7 = self.block7(res2)
+        res3 = self.res3(x1, x7)
+
+        x8 = self.block8(res3)
+        res4 = self.res4(x0, x8)
+
+        x9 = self.block9(res4)
+        y = self.conv15(x9)
+
+        return y
+
+
+class MangaLineExtractor:
+    def __init__(self, model_path: str = "erika.pth", device: str = "cpu"):
+        self.model = res_skip()
+        self.model.load_state_dict(torch.load(model_path))
+
+        self.is_cuda = torch.cuda.is_available() and device == "cuda"
+        if self.is_cuda:
+            self.model.cuda()
+        else:
+            self.model.cpu()
+
+        self.model.eval()
+
+    def predict(self, image):
+        src = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+
+        rows = int(np.ceil(src.shape[0] / 16)) * 16
+        cols = int(np.ceil(src.shape[1] / 16)) * 16
+
+        # manually construct a batch. You can change it based on your usecases.
+        patch = np.ones((1, 1, rows, cols), dtype=np.float32)
+        patch[0, 0, 0 : src.shape[0], 0 : src.shape[1]] = src
+
+        if self.is_cuda:
+            tensor = torch.from_numpy(patch).cuda()
+        else:
+            tensor = torch.from_numpy(patch).cpu()
+
+        y = self.model(tensor)
+
+        yc = y.detach().numpy()[0, 0, :, :]
+        yc[yc > 255] = 255
+        yc[yc < 0] = 0
+        yc = yc / 255.0
+
+        output = yc[0 : src.shape[0], 0 : src.shape[1]]
+        output = cv2.cvtColor(output, cv2.COLOR_GRAY2BGR)
+
+        return output

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+numpy
+opencv-python
+huggingface_hub

setup.py

@@ -0,0 +1,35 @@
+import requests
+import os
+from huggingface_hub import hf_hub_download
+from utils import custom_drive_cache_dir, get_drive
+
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+MANGA_LINE_EXTRACTION_MODEL = "https://github.com/ljsabc/MangaLineExtraction_PyTorch/releases/download/v1/erika.pth"
+ANIME2SKETCH_MODEL = {"REPO_ID": "p1atdev/Anime2Sketch", "FILENAME": "netG.pth"}
+
+
+def download_manga_line_extraction_model():
+    if os.path.exists("./models/erika.pth"):
+        return
+
+
+def download_anime2sketch_model():
+    if os.path.exists("./models/netG.pth"):
+        return
+
+    drive = get_drive("./models/netG.pth")
+    with custom_drive_cache_dir(drive) as cache_dir:
+        hf_hub_download(
+            repo_id=ANIME2SKETCH_MODEL["REPO_ID"],
+            filename=ANIME2SKETCH_MODEL["FILENAME"],
+            local_dir="./models",
+            use_auth_token=HF_TOKEN,
+            local_dir_use_symlinks=False,
+            cache_dir=cache_dir,
+        )
+
+
+def setup():
+    download_manga_line_extraction_model()
+    download_anime2sketch_model()

utils.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+import tempfile
+from contextlib import contextmanager
+import os
+
+
+def get_drive(path: str):
+    path = Path(path).resolve()
+    drive = path.drive
+    root = path.root
+    return drive + root
+
+
+@contextmanager
+def custom_drive_cache_dir(drive: str):
+    drive = Path(drive)
+    base_dir = Path(drive) / "tmp"
+    if not base_dir.exists():
+        os.makedirs(base_dir)
+    print(f"Using {base_dir.resolve()} as cache dir")
+    with tempfile.TemporaryDirectory(dir=base_dir) as tmp_dir:
+        yield tmp_dir