|
1 |
| -# DanBooru IMage Utility functions |
2 |
| - |
3 |
| -import cv2 |
4 |
| -import numpy as np |
5 |
| -from PIL import Image |
6 |
| - |
7 |
| - |
8 |
| -def smart_imread(img, flag=cv2.IMREAD_UNCHANGED): |
9 |
| - if img.endswith(".gif"): |
10 |
| - img = Image.open(img) |
11 |
| - img = img.convert("RGB") |
12 |
| - img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) |
13 |
| - else: |
14 |
| - img = cv2.imread(img, flag) |
15 |
| - return img |
16 |
| - |
17 |
| - |
18 |
| -def smart_24bit(img): |
19 |
| - if img.dtype is np.dtype(np.uint16): |
20 |
| - img = (img / 257).astype(np.uint8) |
21 |
| - |
22 |
| - if len(img.shape) == 2: |
23 |
| - img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) |
24 |
| - elif img.shape[2] == 4: |
25 |
| - trans_mask = img[:, :, 3] == 0 |
26 |
| - img[trans_mask] = [255, 255, 255, 255] |
27 |
| - img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) |
28 |
| - return img |
29 |
| - |
30 |
| - |
31 |
| -def make_square(img, target_size): |
32 |
| - old_size = img.shape[:2] |
33 |
| - desired_size = max(old_size) |
34 |
| - desired_size = max(desired_size, target_size) |
35 |
| - |
36 |
| - delta_w = desired_size - old_size[1] |
37 |
| - delta_h = desired_size - old_size[0] |
38 |
| - top, bottom = delta_h // 2, delta_h - (delta_h // 2) |
39 |
| - left, right = delta_w // 2, delta_w - (delta_w // 2) |
40 |
| - |
41 |
| - color = [255, 255, 255] |
42 |
| - new_im = cv2.copyMakeBorder( |
43 |
| - img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color |
44 |
| - ) |
45 |
| - return new_im |
46 |
| - |
47 |
| - |
48 |
| -def smart_resize(img, size): |
49 |
| - # Assumes the image has already gone through make_square |
50 |
| - if img.shape[0] > size: |
51 |
| - img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA) |
52 |
| - elif img.shape[0] < size: |
53 |
| - img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC) |
54 |
| - return img |
| 1 | +# DanBooru IMage Utility functions |
| 2 | + |
| 3 | +import cv2 |
| 4 | +import numpy as np |
| 5 | +from PIL import Image |
| 6 | + |
| 7 | + |
| 8 | +def smart_imread(img, flag=cv2.IMREAD_UNCHANGED): |
| 9 | + if img.endswith(".gif"): |
| 10 | + img = Image.open(img) |
| 11 | + img = img.convert("RGB") |
| 12 | + img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) |
| 13 | + else: |
| 14 | + img = cv2.imread(img, flag) |
| 15 | + return img |
| 16 | + |
| 17 | + |
| 18 | +def smart_24bit(img): |
| 19 | + if img.dtype is np.dtype(np.uint16): |
| 20 | + img = (img / 257).astype(np.uint8) |
| 21 | + |
| 22 | + if len(img.shape) == 2: |
| 23 | + img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) |
| 24 | + elif img.shape[2] == 4: |
| 25 | + trans_mask = img[:, :, 3] == 0 |
| 26 | + img[trans_mask] = [255, 255, 255, 255] |
| 27 | + img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) |
| 28 | + return img |
| 29 | + |
| 30 | + |
| 31 | +def make_square(img, target_size): |
| 32 | + old_size = img.shape[:2] |
| 33 | + desired_size = max(old_size) |
| 34 | + desired_size = max(desired_size, target_size) |
| 35 | + |
| 36 | + delta_w = desired_size - old_size[1] |
| 37 | + delta_h = desired_size - old_size[0] |
| 38 | + top, bottom = delta_h // 2, delta_h - (delta_h // 2) |
| 39 | + left, right = delta_w // 2, delta_w - (delta_w // 2) |
| 40 | + |
| 41 | + color = [255, 255, 255] |
| 42 | + new_im = cv2.copyMakeBorder( |
| 43 | + img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color |
| 44 | + ) |
| 45 | + return new_im |
| 46 | + |
| 47 | + |
| 48 | +def smart_resize(img, size): |
| 49 | + # Assumes the image has already gone through make_square |
| 50 | + if img.shape[0] > size: |
| 51 | + img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA) |
| 52 | + elif img.shape[0] < size: |
| 53 | + img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC) |
| 54 | + return img |
0 commit comments