File size: 12,428 Bytes
d45e6fb
 
 
 
3503a08
d45e6fb
3503a08
fdce011
d538a3f
3503a08
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d45e6fb
d538a3f
d45e6fb
d538a3f
 
d45e6fb
 
 
d538a3f
 
 
 
 
 
 
d45e6fb
 
d538a3f
 
 
 
d45e6fb
 
 
d538a3f
 
 
 
 
d45e6fb
d538a3f
 
 
 
d45e6fb
d538a3f
 
 
 
 
d45e6fb
d538a3f
 
 
 
d45e6fb
d538a3f
 
 
d45e6fb
d538a3f
 
 
 
 
d45e6fb
d538a3f
 
 
 
 
d45e6fb
 
d538a3f
 
 
 
 
 
 
 
 
 
3503a08
d538a3f
 
 
 
 
 
 
3503a08
d538a3f
3503a08
 
d538a3f
3503a08
d538a3f
 
 
 
3503a08
d538a3f
 
3503a08
 
 
 
 
 
 
 
 
 
 
d538a3f
 
3503a08
d538a3f
 
 
 
3503a08
 
 
d538a3f
 
 
 
 
 
 
 
 
 
3503a08
d538a3f
 
 
 
3503a08
d538a3f
 
 
 
3503a08
d538a3f
 
 
 
 
 
 
 
 
 
 
 
3503a08
d538a3f
3503a08
 
 
 
d538a3f
 
 
 
3503a08
 
d538a3f
 
fdce011
d538a3f
3503a08
d538a3f
 
d45e6fb
3503a08
d45e6fb
 
 
3503a08
d45e6fb
 
d538a3f
 
 
 
3503a08
d538a3f
3503a08
 
d538a3f
 
 
 
3503a08
7cd8e6e
d538a3f
d45e6fb
d538a3f
3503a08
d538a3f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d45e6fb
3503a08
d538a3f
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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
import numpy as np
import torch
import torch.nn as nn
import gradio as gr
from PIL import Image, ImageFilter, ImageOps, ImageChops
import torchvision.transforms as transforms
import os

# 🎨 Dictionary of all filters with their corresponding emojis
FILTERS = {
    "Standard": "📄", "Invert": "⚫⚪", "Blur": "🌫️", "Sharpen": "🔪", "Contour": "🗺️",
    "Detail": "🔍", "EdgeEnhance": "📏", "EdgeEnhanceMore": "📐", "Emboss": "🏞️",
    "FindEdges": "🕵️", "Smooth": "🌊", "SmoothMore": "💧", "Solarize": "☀️",
    "Posterize1": "🖼️1", "Posterize2": "🖼️2", "Posterize3": "🖼️3", "Posterize4": "🖼️4",
    "Equalize": "⚖️", "AutoContrast": "🔧", "Thick1": "💪1", "Thick2": "💪2", "Thick3": "💪3",
    "Thin1": "🏃1", "Thin2": "🏃2", "Thin3": "🏃3", "RedOnWhite": "🔴", "OrangeOnWhite": "🟠",
    "YellowOnWhite": "🟡", "GreenOnWhite": "🟢", "BlueOnWhite": "🔵", "PurpleOnWhite": "🟣",
    "PinkOnWhite": "🌸", "CyanOnWhite": "🩵", "MagentaOnWhite": "🟪", "BrownOnWhite": "🤎",
    "GrayOnWhite": "🩶", "WhiteOnBlack": "⚪", "RedOnBlack": "🔴⚫", "OrangeOnBlack": "🟠⚫",
    "YellowOnBlack": "🟡⚫", "GreenOnBlack": "🟢⚫", "BlueOnBlack": "🔵⚫", "PurpleOnBlack": "🟣⚫",
    "PinkOnBlack": "🌸⚫", "CyanOnBlack": "🩵⚫", "MagentaOnBlack": "🟪⚫", "BrownOnBlack": "🤎⚫",
    "GrayOnBlack": "🩶⚫", "Multiply": "✖️", "Screen": "🖥️", "Overlay": "🔲", "Add": "➕",
    "Subtract": "➖", "Difference": "≠", "Darker": "🌑", "Lighter": "🌕", "SoftLight": "💡",
    "HardLight": "🔦", "Binary": "🌓", "Noise": "❄️"
}

# 🧠 Neural network layers
norm_layer = nn.InstanceNorm2d

# 🧱 Building block for the generator
class ResidualBlock(nn.Module):
    def __init__(self, in_features):
        super(ResidualBlock, self).__init__()
        conv_block = [  nn.ReflectionPad2d(1),
                        nn.Conv2d(in_features, in_features, 3),
                        norm_layer(in_features),
                        nn.ReLU(inplace=True),
                        nn.ReflectionPad2d(1),
                        nn.Conv2d(in_features, in_features, 3),
                        norm_layer(in_features) ]
        self.conv_block = nn.Sequential(*conv_block)

    def forward(self, x):
        return x + self.conv_block(x)

# 🎨 Generator model for creating line drawings
class Generator(nn.Module):
    def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):
        super(Generator, self).__init__()
        # Initial convolution block
        model0 = [   nn.ReflectionPad2d(3),
                    nn.Conv2d(input_nc, 64, 7),
                    norm_layer(64),
                    nn.ReLU(inplace=True) ]
        self.model0 = nn.Sequential(*model0)
        # Downsampling
        model1 = []
        in_features = 64
        out_features = in_features*2
        for _ in range(2):
            model1 += [  nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
                        norm_layer(out_features),
                        nn.ReLU(inplace=True) ]
            in_features = out_features
            out_features = in_features*2
        self.model1 = nn.Sequential(*model1)
        # Residual blocks
        model2 = []
        for _ in range(n_residual_blocks):
            model2 += [ResidualBlock(in_features)]
        self.model2 = nn.Sequential(*model2)
        # Upsampling
        model3 = []
        out_features = in_features//2
        for _ in range(2):
            model3 += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
                        norm_layer(out_features),
                        nn.ReLU(inplace=True) ]
            in_features = out_features
            out_features = in_features//2
        self.model3 = nn.Sequential(*model3)
        # Output layer
        model4 = [  nn.ReflectionPad2d(3),
                    nn.Conv2d(64, output_nc, 7)]
        if sigmoid:
            model4 += [nn.Sigmoid()]
        self.model4 = nn.Sequential(*model4)

    def forward(self, x, cond=None):
        out = self.model0(x)
        out = self.model1(out)
        out = self.model2(out)
        out = self.model3(out)
        out = self.model4(out)
        return out

# 🔧 Load the models
# Make sure you have 'model.pth' and 'model2.pth' in the same directory
try:
    model1 = Generator(3, 1, 3)
    model1.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
    model1.eval()

    model2 = Generator(3, 1, 3)
    model2.load_state_dict(torch.load('model2.pth', map_location=torch.device('cpu')))
    model2.eval()
except FileNotFoundError:
    print("Warning: Model files 'model.pth' or 'model2.pth' not found. The application will not run correctly.")
    model1, model2 = None, None

# ✨ Function to apply the selected filter
def apply_filter(line_img, filter_name, original_img):
    if filter_name == "Standard":
        return line_img
        
    # Convert line drawing to grayscale for most operations
    line_img_l = line_img.convert('L')

    # --- Standard Image Filters ---
    if filter_name == "Invert": return ImageOps.invert(line_img_l)
    if filter_name == "Blur": return line_img.filter(ImageFilter.GaussianBlur(radius=3))
    if filter_name == "Sharpen": return line_img.filter(ImageFilter.SHARPEN)
    if filter_name == "Contour": return line_img_l.filter(ImageFilter.CONTOUR)
    if filter_name == "Detail": return line_img.filter(ImageFilter.DETAIL)
    if filter_name == "EdgeEnhance": return line_img_l.filter(ImageFilter.EDGE_ENHANCE)
    if filter_name == "EdgeEnhanceMore": return line_img_l.filter(ImageFilter.EDGE_ENHANCE_MORE)
    if filter_name == "Emboss": return line_img_l.filter(ImageFilter.EMBOSS)
    if filter_name == "FindEdges": return line_img_l.filter(ImageFilter.FIND_EDGES)
    if filter_name == "Smooth": return line_img.filter(ImageFilter.SMOOTH)
    if filter_name == "SmoothMore": return line_img.filter(ImageFilter.SMOOTH_MORE)

    # --- Tonal Adjustments ---
    if filter_name == "Solarize": return ImageOps.solarize(line_img_l)
    if filter_name == "Posterize1": return ImageOps.posterize(line_img_l, 1)
    if filter_name == "Posterize2": return ImageOps.posterize(line_img_l, 2)
    if filter_name == "Posterize3": return ImageOps.posterize(line_img_l, 3)
    if filter_name == "Posterize4": return ImageOps.posterize(line_img_l, 4)
    if filter_name == "Equalize": return ImageOps.equalize(line_img_l)
    if filter_name == "AutoContrast": return ImageOps.autocontrast(line_img_l)
    if filter_name == "Binary": return line_img_l.convert('1')

    # --- Morphological Operations (Thick/Thin) ---
    if filter_name == "Thick1": return line_img_l.filter(ImageFilter.MinFilter(3))
    if filter_name == "Thick2": return line_img_l.filter(ImageFilter.MinFilter(5))
    if filter_name == "Thick3": return line_img_l.filter(ImageFilter.MinFilter(7))
    if filter_name == "Thin1": return line_img_l.filter(ImageFilter.MaxFilter(3))
    if filter_name == "Thin2": return line_img_l.filter(ImageFilter.MaxFilter(5))
    if filter_name == "Thin3": return line_img_l.filter(ImageFilter.MaxFilter(7))

    # --- Colorization (On White Background) ---
    colors_on_white = {"RedOnWhite": "red", "OrangeOnWhite": "orange", "YellowOnWhite": "yellow", "GreenOnWhite": "green", "BlueOnWhite": "blue", "PurpleOnWhite": "purple", "PinkOnWhite": "pink", "CyanOnWhite": "cyan", "MagentaOnWhite": "magenta", "BrownOnWhite": "brown", "GrayOnWhite": "gray"}
    if filter_name in colors_on_white:
        return ImageOps.colorize(line_img_l, black=colors_on_white[filter_name], white="white")

    # --- Colorization (On Black Background) ---
    colors_on_black = {"WhiteOnBlack": "white", "RedOnBlack": "red", "OrangeOnBlack": "orange", "YellowOnBlack": "yellow", "GreenOnBlack": "green", "BlueOnBlack": "blue", "PurpleOnBlack": "purple", "PinkOnBlack": "pink", "CyanOnBlack": "cyan", "MagentaOnBlack": "magenta", "BrownOnBlack": "brown", "GrayOnBlack": "gray"}
    if filter_name in colors_on_black:
        return ImageOps.colorize(line_img_l, black=colors_on_black[filter_name], white="black")

    # --- Blending Modes with Original Image ---
    line_img_rgb = line_img.convert('RGB')
    if filter_name == "Multiply": return ImageChops.multiply(original_img, line_img_rgb)
    if filter_name == "Screen": return ImageChops.screen(original_img, line_img_rgb)
    if filter_name == "Overlay": return ImageChops.overlay(original_img, line_img_rgb)
    if filter_name == "Add": return ImageChops.add(original_img, line_img_rgb)
    if filter_name == "Subtract": return ImageChops.subtract(original_img, line_img_rgb)
    if filter_name == "Difference": return ImageChops.difference(original_img, line_img_rgb)
    if filter_name == "Darker": return ImageChops.darker(original_img, line_img_rgb)
    if filter_name == "Lighter": return ImageChops.lighter(original_img, line_img_rgb)
    if filter_name == "SoftLight": return ImageChops.soft_light(original_img, line_img_rgb)
    if filter_name == "HardLight": return ImageChops.hard_light(original_img, line_img_rgb)
    
    # --- Texture ---
    if filter_name == "Noise":
        img_array = np.array(line_img_l.convert('L'))
        noise = np.random.randint(-20, 20, img_array.shape, dtype='int16')
        noisy_array = np.clip(img_array.astype('int16') + noise, 0, 255).astype('uint8')
        return Image.fromarray(noisy_array)

    return line_img # Default fallback

# 🖼️ Main function to process the image
def predict(input_img_path, line_style, filter_choice):
    if not model1 or not model2:
        raise gr.Error("Models are not loaded. Please check for 'model.pth' and 'model2.pth'.")
    
    # Extract the filter name from the dropdown choice (e.g., "📄 Standard" -> "Standard")
    filter_name = filter_choice.split(" ", 1)[1]

    original_img = Image.open(input_img_path).convert('RGB')
    original_size = original_img.size

    transform = transforms.Compose([
        transforms.Resize(256, transforms.InterpolationMode.BICUBIC),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    input_tensor = transform(original_img).unsqueeze(0)

    with torch.no_grad():
        if line_style == 'Simple Lines':
            output = model2(input_tensor)
        else: # Complex Lines
            output = model1(input_tensor)
    
    # Convert tensor to low-res PIL image
    line_drawing_low_res = transforms.ToPILImage()(output.squeeze().cpu().clamp(0, 1))
    
    # Resize the line drawing back to the original image size *before* applying filters
    line_drawing_full_res = line_drawing_low_res.resize(original_size, Image.Resampling.BICUBIC)

    # Apply the selected filter
    final_image = apply_filter(line_drawing_full_res, filter_name, original_img)

    return final_image

# 🚀 Setup and launch the Gradio interface
title = "🖌️ Image to Line Art with Creative Filters"
description = "Upload an image, choose a line style (Complex or Simple), and select a filter from the dropdown to transform your picture into unique line art."

# Generate dropdown choices with emojis
filter_choices = [f"{emoji} {name}" for name, emoji in FILTERS.items()]

# Dynamically generate examples from images in the current directory
examples = []
image_dir = '.'
if os.path.exists(image_dir):
    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    if image_files:
        # Pick the first image found and create examples with a few interesting filters
        example_image = image_files[0]
        examples.append([example_image, 'Simple Lines', '🗺️ Contour'])
        examples.append([example_image, 'Complex Lines', '🔵⚫ BlueOnBlack'])
        examples.append([example_image, 'Simple Lines', '✖️ Multiply'])


iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type='filepath', label="Upload Image"),
        gr.Radio(['Complex Lines', 'Simple Lines'], label='Line Style', value='Simple Lines'),
        gr.Dropdown(filter_choices, label="Filter", value=filter_choices[0])
    ],
    outputs=gr.Image(type="pil", label="Filtered Line Art"),
    title=title,
    description=description,
    examples=examples,
    allow_flagging='never'
)

if __name__ == "__main__":
    iface.launch()