ComfyUI-Img2DrawingAssistants/utils.py at main · Isi-dev/ComfyUI-Img2DrawingAssistants · GitHub

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
import numpy as np
import cv2


def convert_to_uint8(image):
    if image.dtype == np.uint8:
        return image
    if image.dtype != np.float32 and image.dtype != np.float64:
        image = image.astype(np.float32)
    if image.max() <= 1.0:
        image = image * 255.0
    image = np.clip(image, 0, 255)
    return image.astype(np.uint8)


def new_bg(data, threshold, penlightness, new_ink_col):
    # threshold = threshold/255
    # mask = (data > threshold).astype(np.float32)
    mask = (data > threshold).astype(np.float32)
    new_ink_color = np.array(new_ink_col, dtype=np.float32)

    if penlightness > 1:
        penlightness = penlightness / 10
        # print(penlightness)
        new_ink_color = new_ink_color + (255 - new_ink_color) * penlightness
    else:
        new_ink_color *= penlightness

    color_data = np.zeros((data.shape[0], data.shape[1], 3), dtype=np.float32)
    color_data[:] = new_ink_color
    color_data = color_data * mask + data * (1-mask)

    return color_data


def contrast_image(data, threshold, scale):
    grayscale = np.mean(data, axis=2)
    mask = grayscale < threshold
    # data[mask] = np.minimum(255, data[mask] * scale)
    data[mask] = data[mask] * scale

    return data


def darken_image(data, threshold, scale):
    # threshold = threshold/255
    mask = data < threshold
    data[mask] = data[mask] * scale

    return data

def new_ink(data, threshold, penlightness, new_ink_col, sketch=False):
    data = data/255
    threshold = threshold/255
    # print(data)
    # print(threshold)
    mask = (data < threshold)
    mask2 = (data >= threshold)
    new_ink_color = np.array(new_ink_col, dtype=np.float32)

    if penlightness > 1:
        # Lighten the color: scale up towards white (255, 255, 255)
        penlightness = penlightness / 10
        # print(penlightness)
        new_ink_color = new_ink_color + (255 - new_ink_color) * penlightness
    else:
        new_ink_color *= penlightness
    # print(new_ink_color)
    color_data = data.astype(np.float32)
    color_data2 = data.astype(np.float32)
    color_data = color_data * (1 - mask) + new_ink_color * data * mask/255
    color_data2 = color_data2 * (1 - mask2) + new_ink_color * data/255
    color_data = color_data + color_data2 * mask2

    return color_data*255


def new_ink2(data, threshold, penlightness, new_ink_col):
    # threshold = threshold/255
    mask = (data > threshold).astype(np.float32)
    new_ink_color = np.array(new_ink_col, dtype=np.float32)

    if penlightness > 1:
        # Lighten the color: scale up towards white (255, 255, 255)
        penlightness = penlightness / 10
        # print(penlightness)
        new_ink_color = new_ink_color + (255 - new_ink_color) * penlightness
    else:
        new_ink_color *= penlightness
    color_data = np.zeros((data.shape[0], data.shape[1], 3), dtype=np.float32)
    color_data[:] = new_ink_color
    color_data = color_data + data * mask

    return color_data

def ensure_grayscale(image):
    if len(image.shape) == 3:
        if image.shape[2] == 4:
            image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    elif len(image.shape) == 2:
        pass  # Already grayscale
    else:
        raise ValueError("Unsupported image format")
    return image


def normalize_image(image):
    # Convert to float if not already
    image = image.astype(np.float32)

    # Normalize to range [0, 1]
    image_min = image.min()
    image_max = image.max()
    # print(f"Image max: {image_max }")

    # Avoid division by zero
    if image_max > image_min:
        normalized_image = (image - image_min) / (image_max - image_min)
    else:
        normalized_image = np.zeros_like(image)
    # print(f"Image max after normalization: {normalized_image.max() }")

    return normalized_image

def remove_noise_from_sketch(image, kn=2, min_size=40, threshold_value=50):
    sketch = ensure_grayscale(image)
    sketch = convert_to_uint8(sketch)
    kernel_size = (kn, kn)

    sketch_inv = cv2.bitwise_not(sketch)

    _, binary_sketch = cv2.threshold(sketch_inv, threshold_value, 255, cv2.THRESH_BINARY)

    kernel = np.ones(kernel_size, np.uint8)
    opened_sketch = cv2.morphologyEx(binary_sketch, cv2.MORPH_OPEN, kernel)

    num_labels, labels_im = cv2.connectedComponents(opened_sketch)

    cleaned_sketch = np.zeros_like(opened_sketch)
    for label in range(1, num_labels):  # Skip the background label (0)
        if np.sum(labels_im == label) >= min_size:
            cleaned_sketch[labels_im == label] = 255

    return cv2.bitwise_not(cleaned_sketch)


def getColor(col):
    color_dict = {
        "white": [255, 255, 255],
        "black": [0, 0, 0],
        "red": [255, 0, 0],
        "lime": [0, 255, 0],
        "blue": [0, 0, 255],
        "yellow": [255, 255, 0],
        "cyan": [0, 255, 255],
        "magenta": [255, 0, 255],
        "silver": [192, 192, 192],
        "gray": [128, 128, 128],
        "maroon": [128, 0, 0],
        "olive": [128, 128, 0],
        "green": [0, 128, 0],
        "purple": [128, 0, 128],
        "teal": [0, 128, 128],
        "navy": [0, 0, 128],
    }

    return color_dict.get(col, [0, 0, 0])