artificialguybr commited on
Commit
d88a29b
1 Parent(s): 6a295ff

Create image_processing.py

Browse files
Files changed (1) hide show
  1. image_processing.py +68 -0
image_processing.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from PIL import Image, ImageDither, ImageQuantize
2
+
3
+ DITHER_METHODS = {
4
+ "None": ImageDither.NONE,
5
+ "Floyd-Steinberg": ImageDither.FLOYDSTEINBERG
6
+ }
7
+
8
+ QUANTIZATION_METHODS = {
9
+ "Median cut": ImageQuantize.MEDIANCUT,
10
+ "Maximum coverage": ImageQuantize.MAXCOVERAGE,
11
+ "Fast octree": ImageQuantize.FASTOCTREE,
12
+ "libimagequant": ImageQuantize.LIBIMAGEQUANT
13
+ }
14
+
15
+ def downscale_image(image: Image, scale: int) -> Image:
16
+ width, height = image.size
17
+ downscaled_image = image.resize((int(width / scale), int(height / scale)), Image.NEAREST)
18
+ return downscaled_image
19
+
20
+ def limit_colors(
21
+ image,
22
+ limit: int=16,
23
+ palette=None,
24
+ palette_colors: int=256,
25
+ quantize: Image.Quantize=Image.Quantize.MEDIANCUT,
26
+ dither: Image.Dither=Image.Dither.NONE,
27
+ use_k_means: bool=False
28
+ ):
29
+ if use_k_means:
30
+ k_means_value = limit
31
+ else:
32
+ k_means_value = 0
33
+
34
+ if palette:
35
+ palette_image = palette
36
+ ppalette = palette.getcolors()
37
+ if ppalette:
38
+ color_palette = palette.quantize(colors=len(list(set(ppalette))))
39
+ else:
40
+ colors = len(palette_image.getcolors()) if palette_image.getcolors() else palette_colors
41
+ color_palette = palette_image.quantize(colors, kmeans=colors)
42
+ else:
43
+ # we need to get palette from image, because
44
+ # dither in quantize doesn't work without it
45
+ # https://pillow.readthedocs.io/en/stable/_modules/PIL/Image.html#Image.quantize
46
+ color_palette = image.quantize(colors=limit, kmeans=k_means_value, method=quantize, dither=Image.Dither.NONE)
47
+
48
+ new_image = image.quantize(palette=color_palette, dither=dither)
49
+
50
+ return new_image
51
+
52
+ def convert_to_grayscale(image):
53
+ new_image = image.convert("L")
54
+ return new_image.convert("RGB")
55
+
56
+ def convert_to_black_and_white(image: Image, threshold: int=128, is_inversed: bool=False):
57
+ if is_inversed:
58
+ apply_threshold = lambda x : 255 if x < threshold else 0
59
+ else:
60
+ apply_threshold = lambda x : 255 if x > threshold else 0
61
+
62
+ black_and_white_image = image.convert('L', dither=Image.Dither.NONE).point(apply_threshold, mode='1')
63
+ return black_and_white_image.convert("RGB")
64
+
65
+ def resize_image(image: Image, size) -> Image:
66
+ width, height = size
67
+ resized_image = image.resize((width, height), Image.NEAREST)
68
+ return resized_image