Skip to content

Commit bddb049

Browse files
developer-accdeveloper-acc
committed
Fixed bug (Couldn't process RGBA and other types of images, now sets image to RGB by default), refactored code
1 parent 82e3779 commit bddb049

5 files changed

Lines changed: 120 additions & 630 deletions

File tree

‎main.py‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import os
22
import json
33
import pathlib
4-
import math
54
from typing import List, Literal, Dict
65

76
import argparse
@@ -81,19 +80,18 @@ def _get_blocks_from_cache(self) -> Dict:
8180
json.dump(blocks, f, indent=4)
8281
return blocks
8382

84-
85-
8683
def convert(self, path: str, output_path: str, show_progress: bool = True) -> None:
87-
if output_path.endswith(".schem"):
84+
if is_video_file(path):
85+
video = mp.VideoFileClip(path)
86+
converted_video = convert_video.process_video_with_pil(video, self.convert_image)
87+
converted_video.write_videofile(output_path, fps=video.fps, logger=None if not show_progress else "bar")
88+
89+
elif output_path.endswith(".schem"):
8890
with Image.open(path, "r") as img:
8991
generate_schematic.create_2d_schematic(
9092
self.get_blocks_2d_matirx(img, show_progress=show_progress),
9193
output_path)
9294

93-
elif is_video_file(path):
94-
video = mp.VideoFileClip(path)
95-
converted_video = convert_video.process_video_with_pil(video, self.convert_image)
96-
converted_video.write_videofile(output_path, fps=video.fps, logger=None if not show_progress else "bar")
9795
else:
9896
with Image.open(path, "r") as img:
9997
converted_image = self.convert_image(img, show_progress=show_progress)
@@ -102,13 +100,15 @@ def convert(self, path: str, output_path: str, show_progress: bool = True) -> No
102100
def preprocess_image(self, image: Image) -> Image:
103101
image_cropper = crop_image.CropImage(image)
104102
cropped_image = image_cropper.crop_to_make_divisible()
105-
103+
if cropped_image.mode != 'RGB':
104+
cropped_image = cropped_image.convert('RGB')
106105
if self.scale_factor > 0 or self.scale_factor < 0:
107106
cropped_image = resize.resize_image(cropped_image, self.scale_factor)
108107

109108
return cropped_image
110109

111110
def get_blocks_2d_matirx(self, image: Image, show_progress: bool = False) -> List[List[str]]:
111+
'''Returns a matrix of strings containing block names.'''
112112
preprocessed_image = self.preprocess_image(image)
113113
width, height = preprocessed_image.size
114114
chunks_x = width // 16
@@ -128,8 +128,8 @@ def get_blocks_2d_matirx(self, image: Image, show_progress: bool = False) -> Lis
128128
lower = upper + 16
129129
chunk = preprocessed_image.crop((left, upper, right, lower))
130130

131-
lowest_block = self.method(chunk)
132-
blocks_matrix[-1].append(lowest_block[0])
131+
closest_block = self.method(chunk)
132+
blocks_matrix[-1].append(closest_block[0])
133133

134134
progress_bar.update(1)
135135

@@ -138,6 +138,7 @@ def get_blocks_2d_matirx(self, image: Image, show_progress: bool = False) -> Lis
138138
return blocks_matrix
139139

140140
def convert_image(self, image: Image, show_progress: bool = False) -> Image:
141+
# TODO: Use get_blocks_2d_matirx to not repeat the code
141142
preprocessed_image = self.preprocess_image(image)
142143

143144
width, height = preprocessed_image.size
@@ -156,8 +157,12 @@ def convert_image(self, image: Image, show_progress: bool = False) -> Image:
156157
lower = upper + 16
157158
chunk = preprocessed_image.crop((left, upper, right, lower))
158159

159-
lowest_block = self.method(chunk)
160-
preprocessed_image.paste(self.blocks_image.crop([self.blocks[lowest_block]["x"], self.blocks[lowest_block]["y"], self.blocks[lowest_block]["x"]+16, self.blocks[lowest_block]["y"]+16]), [left,upper,right,lower])
160+
closest_block = self.method(chunk)
161+
preprocessed_image.paste(
162+
self.blocks_image.crop([self.blocks[closest_block]["x"],
163+
self.blocks[closest_block]["y"], self.blocks[closest_block]["x"]+16,
164+
self.blocks[closest_block]["y"]+16]), [left,upper,right,lower]
165+
)
161166
progress_bar.update(1)
162167

163168
# Close the progress bar
@@ -174,7 +179,7 @@ def main():
174179
# Add the optional arguments
175180
parser.add_argument('--filter', nargs='+', help='Filter options')
176181
parser.add_argument('--scale_factor', type=int, help='Scale factor', default=0)
177-
parser.add_argument('--compression_level', type=int, help='Compression level, greatly improves conversion speed, and loses some information along the way, do not go higher than 20, as it will cause very high memory consumption.', default=16)
182+
parser.add_argument('--compression_level', type=int, help='Compression level, greatly improves conversion speed, and loses some information along the way, do not set higher then 20, as it will cause very high memory consumption.', default=16)
178183
parser.add_argument('--method', type=str,
179184
choices=["abs_diff", "euclidean", "chebyshev_distance", "manhattan_distance", "cosine_similarity", "hamming_distance", "canberra_distance"], help='Method of finding the closest color to block', default="canberra_distance", required=False)
180185
parser.add_argument('--png_atlas_filename', type=str, default=resource_path('minecraft_textures_atlas_blocks.png_0.png'), help='PNG atlas filename')

‎src/calculate_minecraft_blocks_median.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import List, Dict
1+
from typing import Dict
22

33
from PIL import Image, ImageStat
44

‎src/download.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import re
2-
from typing import List, Tuple, Dict
2+
from typing import List, Dict
33
import requests
44

55
from .utils import resource_path

‎src/find_closest.py‎

Lines changed: 99 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
from colorsys import rgb_to_hsv
1+
from typing import Tuple, Callable, Any
22

33
from PIL import Image, ImageStat
44
from scipy.spatial.distance import cosine
55

66
def generate_color_variations(color_dict, max_abs_difference=16):
7+
'''Creates color combinations in given max_abs_difference.'''
78
new_dict = {}
89

910
for rgb_tuple, value in color_dict.items():
@@ -30,110 +31,105 @@ def generate_color_variations(color_dict, max_abs_difference=16):
3031
if total_diff <= max_abs_difference:
3132
new_rgb_tuple = (new_r, new_g, new_b)
3233
new_dict[new_rgb_tuple] = value
33-
3434
return new_dict
3535

3636
class Method:
3737
def __init__(self, blocks, compression_level: int = 16) -> None:
3838
self.compression_level = compression_level
39-
self.caching = dict()
39+
self.cache = dict()
4040
self.blocks = blocks
4141

42+
def add_to_caching(self, median_rgb: Tuple[int, int, int], closest_block: str):
43+
self.cache[median_rgb] = closest_block
44+
new_dict = dict()
45+
new_dict[median_rgb] = closest_block
46+
all_permutations = generate_color_variations(new_dict, self.compression_level)
47+
self.cache.update(all_permutations)
48+
49+
def check_caching(func: Callable[..., Any]):
50+
'''Checks if chunk was already cached, and if so returns cached closest block.'''
51+
def wrapper(self, chunk, *args, **kwargs):
52+
img_median = tuple(ImageStat.Stat(chunk).median)
53+
if img_median in self.cache:
54+
return self.cache[img_median]
55+
else:
56+
return func(self, chunk, *args, **kwargs)
57+
return wrapper
58+
59+
@check_caching
4260
def find_closest_block_rgb_abs_diff(self, chunk: Image) -> str:
4361
'''Calculates the median value of an input image.
4462
Then compares this median to the medians for each block,
4563
and returns the block with the closest match based on the sum of absolute differences between its RGB values and the median of the input image.
4664
If there are multiple blocks with equal minimum difference, it will return the first one encountered.
4765
'''
48-
og_median = tuple(ImageStat.Stat(chunk).median)
49-
og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
50-
if og_median_rgb in self.caching:
51-
return self.caching[og_median_rgb]
52-
else:
53-
rgb_closests_diff = []
54-
for channel in range(3):
55-
min_diff = float('inf')
56-
for block in self.blocks:
57-
diff = abs(og_median_rgb[channel] - self.blocks[block]["median"][channel])
58-
if diff < min_diff:
59-
min_diff = diff
60-
min_diff_block = block
61-
rgb_closests_diff.append(min_diff_block)
62-
63-
lowest_difference = float("inf")
64-
lowest_block = None
65-
for block in rgb_closests_diff:
66-
difference = sum(abs(a - b) for a, b in zip(self.blocks[block]["median"], og_median_rgb))
67-
if difference < lowest_difference:
68-
lowest_difference = difference
69-
lowest_block = block
70-
71-
self.caching[og_median_rgb] = lowest_block
72-
new_dict = dict()
73-
new_dict[og_median_rgb] = lowest_block
74-
all_permutations = generate_color_variations(new_dict, self.compression_level)
75-
self.caching.update(all_permutations)
76-
return lowest_block
66+
img_median = tuple(ImageStat.Stat(chunk).median)
7767

68+
rgb_closests_diff = []
69+
for channel in range(3):
70+
min_diff = float('inf')
71+
for block in self.blocks:
72+
diff = abs(img_median[channel] - self.blocks[block]["median"][channel])
73+
if diff < min_diff:
74+
min_diff = diff
75+
min_diff_block = block
76+
rgb_closests_diff.append(min_diff_block)
77+
78+
lowest_difference = float("inf")
79+
closest_block = None
80+
for block in rgb_closests_diff:
81+
difference = sum(abs(a - b) for a, b in zip(self.blocks[block]["median"], img_median))
82+
if difference < lowest_difference:
83+
lowest_difference = difference
84+
closest_block = block
85+
86+
self.add_to_caching(img_median, closest_block)
87+
return closest_block
88+
89+
@check_caching
7890
def find_closest_block_cosine_similarity(self, chunk: Image) -> str:
7991
'''Calculates the median value of an input image.
8092
Then compares this median to the medians for each block,
8193
and returns the block with the closest match based on the cosine similarity between its RGB values and the median of the input image.
8294
If there are multiple blocks with equal maximum similarity, it will return the first one encountered.
8395
'''
84-
og_median = tuple(ImageStat.Stat(chunk).median)
85-
og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
96+
img_median = tuple(ImageStat.Stat(chunk).median)
97+
98+
closest_block = None
99+
max_similarity = -1
86100

87-
if og_median_rgb in self.caching:
88-
return self.caching[og_median_rgb]
89-
else:
90-
closest_block = None
91-
max_similarity = -1
92-
93-
for block in self.blocks:
94-
block_rgb = self.blocks[block]["median"]
95-
similarity = 1 - cosine(og_median_rgb, block_rgb)
96-
97-
if similarity > max_similarity:
98-
max_similarity = similarity
99-
closest_block = block
101+
for block in self.blocks:
102+
block_rgb = self.blocks[block]["median"]
103+
similarity = 1 - cosine(img_median, block_rgb)
100104

101-
self.caching[og_median_rgb] = closest_block
102-
new_dict = dict()
103-
new_dict[og_median_rgb] = closest_block
104-
all_permutations = generate_color_variations(new_dict, self.compression_level)
105-
self.caching.update(all_permutations)
106-
return closest_block
105+
if similarity > max_similarity:
106+
max_similarity = similarity
107+
closest_block = block
108+
109+
self.add_to_caching(img_median, closest_block)
110+
return closest_block
107111

112+
@check_caching
108113
def find_closest_block_minkowski_distance(self, chunk: Image, p: int=2) -> str:
109114
'''Calculates the median value of an input image.
110115
Then compares this median to the medians for each block,
111116
and returns the block with the closest match based on the Minkowski distance between its RGB values and the median of the input image.
112117
If there are multiple blocks with equal minimum distance, it will return the first one encountered.
113118
'''
114-
og_median = tuple(ImageStat.Stat(chunk).median)
115-
og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
119+
img_median = tuple(ImageStat.Stat(chunk).median)
120+
closest_block = None
121+
min_distance = float('inf')
116122

117-
if og_median_rgb in self.caching:
118-
return self.caching[og_median_rgb]
119-
else:
120-
closest_block = None
121-
min_distance = float('inf')
122-
123-
for block in self.blocks:
124-
block_rgb = self.blocks[block]["median"]
125-
distance = sum(abs(a - b) ** p for a, b in zip(og_median_rgb, block_rgb)) ** (1 / p)
123+
for block in self.blocks:
124+
block_rgb = self.blocks[block]["median"]
125+
distance = sum(abs(a - b) ** p for a, b in zip(img_median, block_rgb)) ** (1 / p)
126126

127-
if distance < min_distance:
128-
min_distance = distance
129-
closest_block = block
127+
if distance < min_distance:
128+
min_distance = distance
129+
closest_block = block
130130

131-
self.caching[og_median_rgb] = closest_block
132-
new_dict = dict()
133-
new_dict[og_median_rgb] = closest_block
134-
all_permutations = generate_color_variations(new_dict, self.compression_level)
135-
self.caching.update(all_permutations)
136-
return closest_block
131+
self.add_to_caching(img_median, closest_block)
132+
return closest_block
137133

138134
def find_closest_block_manhattan_distance(self, chunk: Image) -> str:
139135
return self.find_closest_block_minkowski_distance(chunk, 1)
@@ -147,65 +143,50 @@ def find_closest_block_chebyshev_distance(self, chunk: Image) -> str:
147143
def find_closest_block_taxicab_distance(self, chunk: Image) -> str:
148144
return self.find_closest_block_minkowski_distance(chunk, 4)
149145

146+
@check_caching
150147
def find_closest_block_hamming_distance(self, chunk: Image) -> str:
151148
'''Calculates the median value of an input image.
152149
Then compares this median to the medians for each block,
153150
and returns the block with the closest match based on the Hamming distance between its RGB values and the median of the input image.
154151
If there are multiple blocks with equal minimum distance, it will return the first one encountered.
155152
'''
156-
og_median = tuple(ImageStat.Stat(chunk).median)
157-
og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
153+
img_median = tuple(ImageStat.Stat(chunk).median)
158154

159-
if og_median_rgb in self.caching:
160-
return self.caching[og_median_rgb]
161-
else:
162-
closest_block = None
163-
min_distance = float('inf')
155+
closest_block = None
156+
min_distance = float('inf')
164157

165-
for block in self.blocks:
166-
block_rgb = self.blocks[block]["median"]
167-
distance = sum(a != b for a, b in zip(og_median_rgb, block_rgb))
158+
for block in self.blocks:
159+
block_rgb = self.blocks[block]["median"]
160+
distance = sum(a != b for a, b in zip(img_median, block_rgb))
168161

169-
if distance < min_distance:
170-
min_distance = distance
171-
closest_block = block
162+
if distance < min_distance:
163+
min_distance = distance
164+
closest_block = block
172165

173-
self.caching[og_median_rgb] = closest_block
174-
new_dict = dict()
175-
new_dict[og_median_rgb] = closest_block
176-
all_permutations = generate_color_variations(new_dict, self.compression_level)
177-
self.caching.update(all_permutations)
178-
return closest_block
166+
self.add_to_caching(img_median, closest_block)
167+
return closest_block
179168

169+
@check_caching
180170
def find_closest_block_canberra_distance(self, chunk: Image) -> str:
181171
'''Calculates the median value of an input image.
182172
Then compares this median to the medians for each block,
183173
and returns the block with the closest match based on the Canberra distance between its RGB values and the median of the input image.
184174
If there are multiple blocks with equal minimum distance, it will return the first one encountered.
185175
'''
186-
og_median = tuple(ImageStat.Stat(chunk).median)
187-
og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
188-
189-
if og_median_rgb in self.caching:
190-
return self.caching[og_median_rgb]
191-
else:
192-
closest_block = None
193-
min_distance = float('inf')
194-
195-
for block in self.blocks:
196-
block_rgb = self.blocks[block]["median"]
197-
distance = sum(
198-
abs(a - b) / (abs(a) + abs(b)) if abs(a) + abs(b) != 0 else float('inf')
199-
for a, b in zip(og_median_rgb, block_rgb)
200-
)
201-
202-
if distance < min_distance:
203-
min_distance = distance
204-
closest_block = block
205-
206-
self.caching[og_median_rgb] = closest_block
207-
new_dict = dict()
208-
new_dict[og_median_rgb] = closest_block
209-
all_permutations = generate_color_variations(new_dict, self.compression_level)
210-
self.caching.update(all_permutations)
211-
return closest_block
176+
img_median = tuple(ImageStat.Stat(chunk).median)
177+
closest_block = None
178+
min_distance = float('inf')
179+
180+
for block in self.blocks:
181+
block_rgb = self.blocks[block]["median"]
182+
distance = sum(
183+
abs(a - b) / (abs(a) + abs(b)) if abs(a) + abs(b) != 0 else float('inf')
184+
for a, b in zip(img_median, block_rgb)
185+
)
186+
187+
if distance < min_distance:
188+
min_distance = distance
189+
closest_block = block
190+
191+
self.add_to_caching(img_median, closest_block)
192+
return closest_block

0 commit comments

Comments
 (0)