Fixed bug (Couldn't process RGBA and other types of images, now sets …

…image to RGB by default), refactored code

main.py

@@ -1,7 +1,6 @@
 import os
 import json
 import pathlib
-import math
 from typing import List, Literal, Dict
 
 import argparse
@@ -81,19 +80,18 @@ def _get_blocks_from_cache(self) -> Dict:
 				json.dump(blocks, f, indent=4)
 		return blocks
 
-
-
 	def convert(self, path: str, output_path: str, show_progress: bool = True) -> None:
-		if output_path.endswith(".schem"):
+		if is_video_file(path):
+			video = mp.VideoFileClip(path)
+			converted_video = convert_video.process_video_with_pil(video, self.convert_image)
+			converted_video.write_videofile(output_path, fps=video.fps, logger=None if not show_progress else "bar")
+
+		elif output_path.endswith(".schem"):
 			with Image.open(path, "r") as img:
 				generate_schematic.create_2d_schematic(
 					self.get_blocks_2d_matirx(img, show_progress=show_progress),
 					output_path)
 
-		elif is_video_file(path):
-			video = mp.VideoFileClip(path)
-			converted_video = convert_video.process_video_with_pil(video, self.convert_image)
-			converted_video.write_videofile(output_path, fps=video.fps, logger=None if not show_progress else "bar")
 		else:
 			with Image.open(path, "r") as img:
 				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
 	def preprocess_image(self, image: Image) -> Image:
 		image_cropper = crop_image.CropImage(image)
 		cropped_image = image_cropper.crop_to_make_divisible()
-
+		if cropped_image.mode != 'RGB':
+			cropped_image = cropped_image.convert('RGB')
 		if self.scale_factor > 0 or self.scale_factor < 0:
 			cropped_image = resize.resize_image(cropped_image, self.scale_factor)
 
 		return cropped_image
 
 	def get_blocks_2d_matirx(self, image: Image, show_progress: bool = False) -> List[List[str]]:
+		'''Returns a matrix of strings containing block names.'''
 		preprocessed_image = self.preprocess_image(image)
 		width, height = preprocessed_image.size
 		chunks_x = width // 16
@@ -128,8 +128,8 @@ def get_blocks_2d_matirx(self, image: Image, show_progress: bool = False) -> Lis
 				lower = upper + 16
 				chunk = preprocessed_image.crop((left, upper, right, lower))
 
-				lowest_block = self.method(chunk)
-				blocks_matrix[-1].append(lowest_block[0])
+				closest_block = self.method(chunk)
+				blocks_matrix[-1].append(closest_block[0])
 
 				progress_bar.update(1)
 
@@ -138,6 +138,7 @@ def get_blocks_2d_matirx(self, image: Image, show_progress: bool = False) -> Lis
 		return blocks_matrix
 
 	def convert_image(self, image: Image, show_progress: bool = False) -> Image:
+		# TODO: Use get_blocks_2d_matirx to not repeat the code
 		preprocessed_image = self.preprocess_image(image)
 
 		width, height = preprocessed_image.size
@@ -156,8 +157,12 @@ def convert_image(self, image: Image, show_progress: bool = False) -> Image:
 				lower = upper + 16
 				chunk = preprocessed_image.crop((left, upper, right, lower))
 
-				lowest_block = self.method(chunk)
-				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])
+				closest_block = self.method(chunk)
+				preprocessed_image.paste(
+					self.blocks_image.crop([self.blocks[closest_block]["x"],
+					self.blocks[closest_block]["y"], self.blocks[closest_block]["x"]+16,
+					self.blocks[closest_block]["y"]+16]), [left,upper,right,lower]
+					)
 				progress_bar.update(1)
 
 		# Close the progress bar
@@ -174,7 +179,7 @@ def main():
 	# Add the optional arguments
 	parser.add_argument('--filter', nargs='+', help='Filter options')
 	parser.add_argument('--scale_factor', type=int, help='Scale factor', default=0)
-	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)
+	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)
 	parser.add_argument('--method', type=str,
 		    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)
 	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

@@ -1,4 +1,4 @@
-from typing import List, Dict
+from typing import Dict
 
 from PIL import Image, ImageStat
 

src/download.py

@@ -1,5 +1,5 @@
 import re
-from typing import List, Tuple, Dict
+from typing import List, Dict
 import requests
 
 from .utils import resource_path

src/find_closest.py

@@ -1,9 +1,10 @@
-from colorsys import rgb_to_hsv
+from typing import Tuple, Callable, Any
 
 from PIL import Image, ImageStat
 from scipy.spatial.distance import cosine
 
 def generate_color_variations(color_dict, max_abs_difference=16):
+	'''Creates color combinations in given max_abs_difference.'''
 	new_dict = {}
 
 	for rgb_tuple, value in color_dict.items():
@@ -30,110 +31,105 @@ def generate_color_variations(color_dict, max_abs_difference=16):
 						if total_diff <= max_abs_difference:
 							new_rgb_tuple = (new_r, new_g, new_b)
 							new_dict[new_rgb_tuple] = value
-
 	return new_dict
 
 class Method:
 	def __init__(self, blocks, compression_level: int = 16) -> None:
 		self.compression_level = compression_level
-		self.caching = dict()
+		self.cache = dict()
 		self.blocks = blocks
 
+	def add_to_caching(self, median_rgb: Tuple[int, int, int], closest_block: str):
+		self.cache[median_rgb] = closest_block
+		new_dict = dict()
+		new_dict[median_rgb] = closest_block
+		all_permutations = generate_color_variations(new_dict, self.compression_level)
+		self.cache.update(all_permutations)
+
+	def check_caching(func: Callable[..., Any]):
+		'''Checks if chunk was already cached, and if so returns cached closest block.'''
+		def wrapper(self, chunk, *args, **kwargs):
+			img_median = tuple(ImageStat.Stat(chunk).median)
+			if img_median in self.cache:
+				return self.cache[img_median]
+			else:
+				return func(self, chunk, *args, **kwargs)
+		return wrapper
+
+	@check_caching
 	def find_closest_block_rgb_abs_diff(self, chunk: Image) -> str:
 		'''Calculates the median value of an input image.
 		Then compares this median to the medians for each block,
 		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.
 		If there are multiple blocks with equal minimum difference, it will return the first one encountered.
 		'''
-		og_median = tuple(ImageStat.Stat(chunk).median)
-		og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
-		if og_median_rgb in self.caching:
-			return self.caching[og_median_rgb]
-		else:
-			rgb_closests_diff = []
-			for channel in range(3):
-				min_diff = float('inf')
-				for block in self.blocks:
-					diff = abs(og_median_rgb[channel] - self.blocks[block]["median"][channel])
-					if diff < min_diff:
-						min_diff = diff
-						min_diff_block = block
-				rgb_closests_diff.append(min_diff_block)
-
-			lowest_difference = float("inf")
-			lowest_block = None
-			for block in rgb_closests_diff:
-				difference = sum(abs(a - b) for a, b in zip(self.blocks[block]["median"], og_median_rgb))
-				if difference < lowest_difference:
-					lowest_difference = difference
-					lowest_block = block
-
-			self.caching[og_median_rgb] = lowest_block
-			new_dict = dict()
-			new_dict[og_median_rgb] = lowest_block
-			all_permutations = generate_color_variations(new_dict, self.compression_level)
-			self.caching.update(all_permutations)
-			return lowest_block
+		img_median = tuple(ImageStat.Stat(chunk).median)
 
+		rgb_closests_diff = []
+		for channel in range(3):
+			min_diff = float('inf')
+			for block in self.blocks:
+				diff = abs(img_median[channel] - self.blocks[block]["median"][channel])
+				if diff < min_diff:
+					min_diff = diff
+					min_diff_block = block
+			rgb_closests_diff.append(min_diff_block)
+
+		lowest_difference = float("inf")
+		closest_block = None
+		for block in rgb_closests_diff:
+			difference = sum(abs(a - b) for a, b in zip(self.blocks[block]["median"], img_median))
+			if difference < lowest_difference:
+				lowest_difference = difference
+				closest_block = block
+
+		self.add_to_caching(img_median, closest_block)
+		return closest_block
+
+	@check_caching
 	def find_closest_block_cosine_similarity(self, chunk: Image) -> str:
 		'''Calculates the median value of an input image.
 		Then compares this median to the medians for each block,
 		and returns the block with the closest match based on the cosine similarity between its RGB values and the median of the input image.
 		If there are multiple blocks with equal maximum similarity, it will return the first one encountered.
 		'''
-		og_median = tuple(ImageStat.Stat(chunk).median)
-		og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
+		img_median = tuple(ImageStat.Stat(chunk).median)
+
+		closest_block = None
+		max_similarity = -1
 
-		if og_median_rgb in self.caching:
-			return self.caching[og_median_rgb]
-		else:
-			closest_block = None
-			max_similarity = -1
-
-			for block in self.blocks:
-				block_rgb = self.blocks[block]["median"]
-				similarity = 1 - cosine(og_median_rgb, block_rgb)
-
-				if similarity > max_similarity:
-					max_similarity = similarity
-					closest_block = block
+		for block in self.blocks:
+			block_rgb = self.blocks[block]["median"]
+			similarity = 1 - cosine(img_median, block_rgb)
 
-			self.caching[og_median_rgb] = closest_block
-			new_dict = dict()
-			new_dict[og_median_rgb] = closest_block
-			all_permutations = generate_color_variations(new_dict, self.compression_level)
-			self.caching.update(all_permutations)
-			return closest_block
+			if similarity > max_similarity:
+				max_similarity = similarity
+				closest_block = block
+
+		self.add_to_caching(img_median, closest_block)
+		return closest_block
 
+	@check_caching
 	def find_closest_block_minkowski_distance(self, chunk: Image, p: int=2) -> str:
 		'''Calculates the median value of an input image.
 		Then compares this median to the medians for each block,
 		and returns the block with the closest match based on the Minkowski distance between its RGB values and the median of the input image.
 		If there are multiple blocks with equal minimum distance, it will return the first one encountered.
 		'''
-		og_median = tuple(ImageStat.Stat(chunk).median)
-		og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
+		img_median = tuple(ImageStat.Stat(chunk).median)
+		closest_block = None
+		min_distance = float('inf')
 
-		if og_median_rgb in self.caching:
-			return self.caching[og_median_rgb]
-		else:
-			closest_block = None
-			min_distance = float('inf')
-
-			for block in self.blocks:
-				block_rgb = self.blocks[block]["median"]
-				distance = sum(abs(a - b) ** p for a, b in zip(og_median_rgb, block_rgb)) ** (1 / p)
+		for block in self.blocks:
+			block_rgb = self.blocks[block]["median"]
+			distance = sum(abs(a - b) ** p for a, b in zip(img_median, block_rgb)) ** (1 / p)
 
-				if distance < min_distance:
-					min_distance = distance
-					closest_block = block
+			if distance < min_distance:
+				min_distance = distance
+				closest_block = block
 
-			self.caching[og_median_rgb] = closest_block
-			new_dict = dict()
-			new_dict[og_median_rgb] = closest_block
-			all_permutations = generate_color_variations(new_dict, self.compression_level)
-			self.caching.update(all_permutations)
-			return closest_block
+		self.add_to_caching(img_median, closest_block)
+		return closest_block
 
 	def find_closest_block_manhattan_distance(self, chunk: Image) -> str:
 		return self.find_closest_block_minkowski_distance(chunk, 1)
@@ -147,65 +143,50 @@ def find_closest_block_chebyshev_distance(self, chunk: Image) -> str:
 	def find_closest_block_taxicab_distance(self, chunk: Image) -> str:
 		return self.find_closest_block_minkowski_distance(chunk, 4)
 
+	@check_caching
 	def find_closest_block_hamming_distance(self, chunk: Image) -> str:
 		'''Calculates the median value of an input image.
 		Then compares this median to the medians for each block,
 		and returns the block with the closest match based on the Hamming distance between its RGB values and the median of the input image.
 		If there are multiple blocks with equal minimum distance, it will return the first one encountered.
 		'''
-		og_median = tuple(ImageStat.Stat(chunk).median)
-		og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
+		img_median = tuple(ImageStat.Stat(chunk).median)
 
-		if og_median_rgb in self.caching:
-			return self.caching[og_median_rgb]
-		else:
-			closest_block = None
-			min_distance = float('inf')
+		closest_block = None
+		min_distance = float('inf')
 
-			for block in self.blocks:
-				block_rgb = self.blocks[block]["median"]
-				distance = sum(a != b for a, b in zip(og_median_rgb, block_rgb))
+		for block in self.blocks:
+			block_rgb = self.blocks[block]["median"]
+			distance = sum(a != b for a, b in zip(img_median, block_rgb))
 
-				if distance < min_distance:
-					min_distance = distance
-					closest_block = block
+			if distance < min_distance:
+				min_distance = distance
+				closest_block = block
 
-			self.caching[og_median_rgb] = closest_block
-			new_dict = dict()
-			new_dict[og_median_rgb] = closest_block
-			all_permutations = generate_color_variations(new_dict, self.compression_level)
-			self.caching.update(all_permutations)
-			return closest_block
+		self.add_to_caching(img_median, closest_block)
+		return closest_block
 
+	@check_caching
 	def find_closest_block_canberra_distance(self, chunk: Image) -> str:
 		'''Calculates the median value of an input image.
 		Then compares this median to the medians for each block,
 		and returns the block with the closest match based on the Canberra distance between its RGB values and the median of the input image.
 		If there are multiple blocks with equal minimum distance, it will return the first one encountered.
 		'''
-		og_median = tuple(ImageStat.Stat(chunk).median)
-		og_median_rgb = tuple([og_median[0], og_median[1], og_median[2]])
-
-		if og_median_rgb in self.caching:
-			return self.caching[og_median_rgb]
-		else:
-			closest_block = None
-			min_distance = float('inf')
-
-			for block in self.blocks:
-				block_rgb = self.blocks[block]["median"]
-				distance = sum(
-					abs(a - b) / (abs(a) + abs(b)) if abs(a) + abs(b) != 0 else float('inf')
-					for a, b in zip(og_median_rgb, block_rgb)
-				)
-
-				if distance < min_distance:
-					min_distance = distance
-					closest_block = block
-
-			self.caching[og_median_rgb] = closest_block
-			new_dict = dict()
-			new_dict[og_median_rgb] = closest_block
-			all_permutations = generate_color_variations(new_dict, self.compression_level)
-			self.caching.update(all_permutations)
-			return closest_block
+		img_median = tuple(ImageStat.Stat(chunk).median)
+		closest_block = None
+		min_distance = float('inf')
+
+		for block in self.blocks:
+			block_rgb = self.blocks[block]["median"]
+			distance = sum(
+				abs(a - b) / (abs(a) + abs(b)) if abs(a) + abs(b) != 0 else float('inf')
+				for a, b in zip(img_median, block_rgb)
+			)
+
+			if distance < min_distance:
+				min_distance = distance
+				closest_block = block
+
+		self.add_to_caching(img_median, closest_block)
+		return closest_block