Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
sketch_simplification.py 4.34 KB
Newer Older
mashun1's avatar
i2vgen-xl  
mashun1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114 
r"""PyTorch re-implementation adapted from the Lua code in ``https://github.com/bobbens/sketch_simplification''.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# from canvas import DOWNLOAD_TO_CACHE
from artist import DOWNLOAD_TO_CACHE

__all__ = ['SketchSimplification', 'sketch_simplification_gan', 'sketch_simplification_mse',
           'sketch_to_pencil_v1', 'sketch_to_pencil_v2']

class SketchSimplification(nn.Module):
    r"""NOTE:
        1. Input image should has only one gray channel.
        2. Input image size should be divisible by 8.
        3. Sketch in the input/output image is in dark color while background in light color.
    """
    def __init__(self, mean, std):
        assert isinstance(mean, float) and isinstance(std, float)
        super(SketchSimplification, self).__init__()
        self.mean = mean
        self.std = std

        # layers
        self.layers = nn.Sequential(
            nn.Conv2d(1, 48, 5, 2, 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(48, 128, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, 2, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, 2, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(1024, 512, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 256, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(256, 256, 4, 2, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 128, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(128, 128, 4, 2, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 48, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(48, 48, 4, 2, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(48, 24, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(24, 1, 3, 1, 1),
            nn.Sigmoid())
    
    def forward(self, x):
        r"""x: [B, 1, H, W] within range [0, 1]. Sketch pixels in dark color.
        """
        x = (x - self.mean) / self.std
        return self.layers(x)

def sketch_simplification_gan(pretrained=False):
    model = SketchSimplification(mean=0.9664114577640158, std=0.0858381272736797)
    if pretrained:
        # model.load_state_dict(torch.load(
        #     DOWNLOAD_TO_CACHE('models/sketch_simplification/sketch_simplification_gan.pth'),
        #     map_location='cpu'))
        model.load_state_dict(torch.load(
            DOWNLOAD_TO_CACHE('VideoComposer/Hangjie/models/sketch_simplification/sketch_simplification_gan.pth'),
            map_location='cpu'))
    return model

def sketch_simplification_mse(pretrained=False):
    model = SketchSimplification(mean=0.9664423107454593, std=0.08583666033640507)
    if pretrained:
        model.load_state_dict(torch.load(
            DOWNLOAD_TO_CACHE('models/sketch_simplification/sketch_simplification_mse.pth'),
            map_location='cpu'))
    return model

def sketch_to_pencil_v1(pretrained=False):
    model = SketchSimplification(mean=0.9817833515894078, std=0.0925009022585048)
    if pretrained:
        model.load_state_dict(torch.load(
            DOWNLOAD_TO_CACHE('models/sketch_simplification/sketch_to_pencil_v1.pth'),
            map_location='cpu'))
    return model

def sketch_to_pencil_v2(pretrained=False):
    model = SketchSimplification(mean=0.9851298627337799, std=0.07418377454883571)
    if pretrained:
        model.load_state_dict(torch.load(
            DOWNLOAD_TO_CACHE('models/sketch_simplification/sketch_to_pencil_v2.pth'),
            map_location='cpu'))
    return model