646 lines
27 KiB
Python
646 lines
27 KiB
Python
import numpy as np
|
|
import qrcode
|
|
from qrcode.image.styles.moduledrawers import (GappedSquareModuleDrawer,
|
|
CircleModuleDrawer,
|
|
RoundedModuleDrawer,
|
|
VerticalBarsDrawer,
|
|
HorizontalBarsDrawer)
|
|
from qrcode.image.styledpil import StyledPilImage
|
|
from qrcode.image.styles.colormasks import SolidFillColorMask
|
|
from PIL import Image
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class QRBase:
|
|
def __init__(self):
|
|
self.text = ""
|
|
self.fill = None
|
|
self.back = None
|
|
|
|
FUNCTION = "generate_qr"
|
|
CATEGORY = "ComfyQR"
|
|
|
|
def _get_error_correction_constant(self, error_correction_string):
|
|
if error_correction_string == "Low":
|
|
return qrcode.constants.ERROR_CORRECT_L
|
|
if error_correction_string == "Medium":
|
|
return qrcode.constants.ERROR_CORRECT_M
|
|
if error_correction_string == "Quartile":
|
|
return qrcode.constants.ERROR_CORRECT_Q
|
|
return qrcode.constants.ERROR_CORRECT_H
|
|
|
|
def _img_to_tensor(self, img):
|
|
out_image = np.array(img, dtype=np.uint8).astype(np.float32) / 255
|
|
return torch.from_numpy(out_image).unsqueeze(0)
|
|
|
|
def _make_qr(self, qr, fill_hexcolor, back_hexcolor, module_drawer):
|
|
self.fill = self._parse_hexcolor_string(fill_hexcolor, "fill_hexcolor")
|
|
self.back = self._parse_hexcolor_string(back_hexcolor, "back_hexcolor")
|
|
qr.make(fit=True)
|
|
if module_drawer == "Square":
|
|
# Keeps using Square QR generation the old way for faster speeds.
|
|
return qr.make_image(fill_color=self.fill, back_color=self.back)
|
|
color_mask = SolidFillColorMask(back_color=self.back,
|
|
front_color=self.fill)
|
|
module_drawing_method = self._select_module_drawer(module_drawer)
|
|
return qr.make_image(image_factory=StyledPilImage,
|
|
color_mask=color_mask,
|
|
module_drawer=module_drawing_method)
|
|
|
|
def _parse_hexcolor_string(self, s, parameter):
|
|
if s.startswith("#"):
|
|
s = s[1:]
|
|
if len(s) == 3:
|
|
rgb = (c + c for c in s)
|
|
elif len(s) == 6:
|
|
rgb = (s[i] + s[i+1] for i in range(0, 6, 2))
|
|
else:
|
|
raise ValueError(f"{parameter} must be 3 or 6 characters long")
|
|
try:
|
|
return tuple(int(channel, 16) for channel in rgb)
|
|
except ValueError:
|
|
raise ValueError(f"{parameter} contains invalid hexadecimal "
|
|
f"characters")
|
|
|
|
def _validate_qr_size(self, size, max_size):
|
|
if size > max_size:
|
|
raise RuntimeError(f"QR dimensions of {size} exceed max size of "
|
|
f"{max_size}.")
|
|
|
|
def _select_module_drawer(self, module_drawer_string):
|
|
"""Square is not included in the results, for a speed optimization
|
|
applying color masks. Current version of python-qr code suffers a
|
|
slowdown when using custom colors combined with custom module drawers.
|
|
By bypassing square QRs, non standard colors will load faster."""
|
|
if module_drawer_string == "Gapped square":
|
|
return GappedSquareModuleDrawer()
|
|
if module_drawer_string == "Circle":
|
|
return CircleModuleDrawer()
|
|
if module_drawer_string == "Rounded":
|
|
return RoundedModuleDrawer()
|
|
if module_drawer_string == "Vertical bars":
|
|
return VerticalBarsDrawer()
|
|
if module_drawer_string == "Horizontal bars":
|
|
return HorizontalBarsDrawer()
|
|
raise ValueError(f"Module drawing method of {module_drawer_string} "
|
|
f"not supported")
|
|
|
|
def update_text(self, protocol, text):
|
|
"""This function takes input from a text box and a chosen internet
|
|
protocol and stores a full address within an instance variable.
|
|
Backslashes will invalidate text box input and this acts as a
|
|
workaround to be able to use them when required in QR strings.
|
|
|
|
Args:
|
|
protocol: A categorical variable of one of the available internet
|
|
protocols.
|
|
text: The input from the text box.
|
|
"""
|
|
if protocol == "Https":
|
|
prefix = "https://"
|
|
elif protocol == "Http":
|
|
prefix = "http://"
|
|
elif protocol == "None":
|
|
prefix = ""
|
|
self.text = prefix + text
|
|
|
|
|
|
class QRByImageSize(QRBase):
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"protocol": (["Http", "Https", "None"], {"default": "Https"}),
|
|
"text": ("STRING", {"multiline": True}),
|
|
"image_size": ("INT", {"default": 512,
|
|
"min": 64,
|
|
"max": 4096,
|
|
"step": 64}),
|
|
"fill_hexcolor": ("STRING", {"multiline": False,
|
|
"default": "#000000"}),
|
|
"back_hexcolor": ("STRING", {"multiline": False,
|
|
"default": "#FFFFFF"}),
|
|
"error_correction": (["Low", "Medium", "Quartile", "High"],
|
|
{"default": "High"}),
|
|
"border": ("INT", {"default": 1,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 1}),
|
|
"resampling": (["Bicubic",
|
|
"Bilinear",
|
|
"Box",
|
|
"Hamming",
|
|
"Lanczos",
|
|
"Nearest"
|
|
], {"default": "Nearest"}),
|
|
"module_drawer": (["Square",
|
|
"Gapped square",
|
|
"Circle",
|
|
"Rounded",
|
|
"Vertical bars",
|
|
"Horizontal bars"
|
|
], {"default": "Square"})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT")
|
|
RETURN_NAMES = ("QR_CODE", "QR_VERSION")
|
|
|
|
def _select_resampling_method(self, resampling_string):
|
|
if resampling_string == "Nearest":
|
|
return Image.NEAREST
|
|
if resampling_string == "Bicubic":
|
|
return Image.BICUBIC
|
|
if resampling_string == "Bilinear":
|
|
return Image.BILINEAR
|
|
if resampling_string == "Lanczos":
|
|
return Image.LANCZOS
|
|
if resampling_string == "Box":
|
|
return Image.BOX
|
|
if resampling_string == "Hamming":
|
|
return Image.HAMMING
|
|
raise ValueError(f"Resampling method of {resampling_string} not "
|
|
f"supported")
|
|
|
|
def generate_qr(
|
|
self,
|
|
protocol,
|
|
text,
|
|
image_size,
|
|
fill_hexcolor,
|
|
back_hexcolor,
|
|
error_correction,
|
|
border,
|
|
resampling,
|
|
module_drawer
|
|
):
|
|
resampling_method = self._select_resampling_method(resampling)
|
|
error_level = self._get_error_correction_constant(error_correction)
|
|
self.update_text(protocol, text)
|
|
qr = qrcode.QRCode(
|
|
error_correction=error_level,
|
|
box_size=16,
|
|
border=border)
|
|
qr.add_data(self.text)
|
|
img = self._make_qr(qr, fill_hexcolor, back_hexcolor, module_drawer)
|
|
img = img.resize((image_size, image_size), resample=resampling_method)
|
|
return (self._img_to_tensor(img), qr.version)
|
|
|
|
|
|
class QRByModuleSize(QRBase):
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"protocol": (["Http", "Https", "None"], {"default": "Https"}),
|
|
"text": ("STRING", {"multiline": True}),
|
|
"module_size": ("INT", {"default": 16,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1}),
|
|
"max_image_size": ("INT", {"default": 512,
|
|
"min": 64,
|
|
"max": 4096,
|
|
"step": 64}),
|
|
"fill_hexcolor": ("STRING", {"multiline": False,
|
|
"default": "#000000"}),
|
|
"back_hexcolor": ("STRING", {"multiline": False,
|
|
"default": "#FFFFFF"}),
|
|
"error_correction": (["Low", "Medium", "Quartile", "High"],
|
|
{"default": "High"}),
|
|
"border": ("INT", {"default": 1,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 1}),
|
|
"module_drawer": (["Square",
|
|
"Gapped square",
|
|
"Circle",
|
|
"Rounded",
|
|
"Vertical bars",
|
|
"Horizontal bars"
|
|
], {"default": "Square"})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT")
|
|
RETURN_NAMES = ("QR_CODE", "QR_VERSION", "IMAGE_SIZE")
|
|
|
|
def generate_qr(
|
|
self,
|
|
protocol,
|
|
text,
|
|
module_size,
|
|
max_image_size,
|
|
fill_hexcolor,
|
|
back_hexcolor,
|
|
error_correction,
|
|
border,
|
|
module_drawer
|
|
):
|
|
self.update_text(protocol, text)
|
|
error_level = self._get_error_correction_constant(error_correction)
|
|
qr = qrcode.QRCode(
|
|
error_correction=error_level,
|
|
box_size=module_size,
|
|
border=border)
|
|
qr.add_data(self.text)
|
|
img = self._make_qr(qr, fill_hexcolor, back_hexcolor, module_drawer)
|
|
self._validate_qr_size(img.pixel_size, max_image_size)
|
|
return (self._img_to_tensor(img), qr.version, img.pixel_size)
|
|
|
|
|
|
class QRByModuleSizeSplitFunctionPatterns(QRBase):
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"protocol": (["Http", "Https", "None"], {"default": "Https"}),
|
|
"text": ("STRING", {"multiline": True}),
|
|
"module_size": ("INT", {"default": 16,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1}),
|
|
"max_image_size": ("INT", {"default": 512,
|
|
"min": 64,
|
|
"max": 4096,
|
|
"step": 64}),
|
|
"fill_hexcolor": ("STRING", {"multiline": False,
|
|
"default": "#000000"}),
|
|
"back_hexcolor": ("STRING", {"multiline": False,
|
|
"default": "#FFFFFF"}),
|
|
"error_correction": (["Low", "Medium", "Quartile", "High"],
|
|
{"default": "High"}),
|
|
"border": ("INT", {"default": 1,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 1}),
|
|
"module_drawer": (["Square",
|
|
"Gapped square",
|
|
"Circle",
|
|
"Rounded",
|
|
"Vertical bars",
|
|
"Horizontal bars"
|
|
], {"default": "Square"})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "MASK", "INT", "INT")
|
|
RETURN_NAMES = ("QR_CODE",
|
|
"MODULE_LAYER",
|
|
"FINDER_LAYER",
|
|
"FINDER_MASK",
|
|
"QR_VERSION",
|
|
"IMAGE_SIZE")
|
|
|
|
def _generate_finder_pattern_ranges(self, module_size, border_size):
|
|
outer = module_size * border_size
|
|
inner = 7 * module_size + outer
|
|
# Alternate behavior is required to prevent bugs from 0 border_size.
|
|
far_outer = -outer if border_size else None
|
|
return [
|
|
(outer, inner, outer, inner),
|
|
(outer, inner, -inner, far_outer),
|
|
(-inner, far_outer, outer, inner)
|
|
]
|
|
|
|
def _generate_finder_pattern_mask(self,
|
|
pixel_size,
|
|
module_size,
|
|
border_size):
|
|
mask = np.zeros((pixel_size, pixel_size), dtype=bool)
|
|
for (x_min,
|
|
x_max,
|
|
y_min,
|
|
y_max) in self._generate_finder_pattern_ranges(module_size,
|
|
border_size):
|
|
mask[y_min:y_max, x_min:x_max] = True
|
|
return mask
|
|
|
|
def _apply_fill_to_mask(self, img, mask):
|
|
array = np.array(img).copy()
|
|
indices = np.nonzero(mask)
|
|
array[indices[0], indices[1], :] = self.back
|
|
return Image.fromarray(array)
|
|
|
|
def _mask_to_tensor(self, mask):
|
|
out_image = mask.astype(np.float32)
|
|
return torch.from_numpy(out_image).unsqueeze(0)
|
|
|
|
def generate_qr(
|
|
self,
|
|
protocol,
|
|
text,
|
|
module_size,
|
|
max_image_size,
|
|
fill_hexcolor,
|
|
back_hexcolor,
|
|
error_correction,
|
|
border,
|
|
module_drawer
|
|
):
|
|
self.update_text(protocol, text)
|
|
error_level = self._get_error_correction_constant(error_correction)
|
|
qr = qrcode.QRCode(
|
|
error_correction=error_level,
|
|
box_size=module_size,
|
|
border=border)
|
|
qr.add_data(self.text)
|
|
img = self._make_qr(qr, fill_hexcolor, back_hexcolor, module_drawer)
|
|
pixel_size = img.pixel_size
|
|
self._validate_qr_size(pixel_size, max_image_size)
|
|
mask = self._generate_finder_pattern_mask(pixel_size,
|
|
module_size,
|
|
border)
|
|
module_image = self._apply_fill_to_mask(img, mask)
|
|
function_image = self._apply_fill_to_mask(img, ~mask)
|
|
return (
|
|
self._img_to_tensor(img),
|
|
self._img_to_tensor(module_image),
|
|
self._img_to_tensor(function_image),
|
|
self._mask_to_tensor(mask),
|
|
qr.version,
|
|
pixel_size,
|
|
)
|
|
|
|
|
|
class QRErrorMasker:
|
|
def __init__(self):
|
|
self.module_size = None
|
|
self.canvas_shape = None
|
|
self.qr_bounds = None
|
|
|
|
FUNCTION = "find_qr_errors"
|
|
CATEGORY = "ComfyQR"
|
|
RETURN_TYPES = ("MASK", "FLOAT", "FLOAT", "FLOAT")
|
|
RETURN_NAMES = ("QR_ERROR_MASK", "PERCENT_ERROR", "CORRELATION", "RMSE")
|
|
OUTPUT_IS_LIST = (False, True, True, True)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source_qr": ("IMAGE",),
|
|
"modified_qr": ("IMAGE",),
|
|
"module_size": ("INT", {"default": 16,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1}),
|
|
"grayscale_method": (["mean", "luminance"],
|
|
{"default": "luminance"}),
|
|
"aggregate_method": (["mean",], {"default": "mean"}),
|
|
"evaluate": (["full_qr", "module_pattern", "finder_pattern"],
|
|
{"default": "module_pattern"}),
|
|
"error_difficulty": ("FLOAT", {"default": 0,
|
|
"min": 0,
|
|
"max": 1,
|
|
"step": .01}),
|
|
"inverted_pattern": ("BOOLEAN", {"default": False}),
|
|
"gamma": ("FLOAT", {"default": 2.2,
|
|
"min": .1,
|
|
"max": 2.8,
|
|
"step": .1}),
|
|
},
|
|
}
|
|
|
|
def _get_qr_bounds(self, tensor, invert):
|
|
module_color = 1.0 if invert else 0.0
|
|
module_pixels = (tensor == module_color)
|
|
indices = torch.nonzero(module_pixels, as_tuple=True)
|
|
# The viewer patterns will guarentee a module pixel in the upper left
|
|
# The bottom right does not have that guarentee so max is used.
|
|
return (indices[0][0],
|
|
indices[0].max() + 1,
|
|
indices[1][0], indices[1].max() + 1)
|
|
|
|
def _extract_pattern_from_bounds(self, tensor):
|
|
return tensor[self.qr_bounds[0]: self.qr_bounds[1],
|
|
self.qr_bounds[2]: self.qr_bounds[3]]
|
|
|
|
def _trim_to_qr_area(self, source_qr, modified_qr, inverted_pattern):
|
|
self.qr_bounds = self._get_qr_bounds(source_qr, inverted_pattern)
|
|
self._check_bounds_and_module_size()
|
|
source_qr = self._extract_pattern_from_bounds(source_qr)
|
|
modified_qr = self._extract_pattern_from_bounds(modified_qr)
|
|
return source_qr, modified_qr
|
|
|
|
def _reshape_tensor_to_modules(self, tensor):
|
|
if len(tensor.shape) != 2:
|
|
raise RuntimeError("Module reshaping requires a 2 dimensional "
|
|
"array.")
|
|
length = tensor.shape[0] // self.module_size
|
|
reshaped_tensor = tensor.view(length,
|
|
self.module_size,
|
|
length,
|
|
self.module_size)
|
|
rehaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()
|
|
return rehaped_tensor.view(length, length, self.module_size ** 2)
|
|
|
|
def _check_bounds_and_module_size(self):
|
|
height = self.qr_bounds[1] - self.qr_bounds[0]
|
|
width = self.qr_bounds[3] - self.qr_bounds[2]
|
|
color_warning = "Make sure that qr_fill and back colors have exact "
|
|
"#FFFFFFF and #000000 values (and that module color values do not "
|
|
"occur outside the QR) and invert is set correctly."
|
|
if width != height:
|
|
raise RuntimeError(f"Source QR dimensions are {width} x {height}. "
|
|
f"They must be a perfect square. "
|
|
f"{color_warning}")
|
|
if width % self.module_size:
|
|
raise RuntimeError(f"QR width of {width} does not fit module_size "
|
|
f"of {self.module_size}. It must be perfectly "
|
|
f"divisible. {color_warning}")
|
|
|
|
def _check_equal_shape(self, source_qr, modified_qr):
|
|
if source_qr.shape != modified_qr.shape:
|
|
raise ValueError("Source and modified QR must have the same batch "
|
|
"size and dimensions.")
|
|
|
|
def _squeeze_by_mean(self, tensor):
|
|
return torch.mean(tensor, dim=-1)
|
|
|
|
def _gamma_expansion(self, tensor, gamma):
|
|
if gamma == 1:
|
|
return tensor
|
|
if gamma == 2.2:
|
|
return torch.where(tensor <= 0.04045,
|
|
tensor / 12.92,
|
|
(tensor + 0.055) / 1.055) ** 2.4
|
|
return tensor ** gamma
|
|
|
|
def _gamma_compression(self, tensor, gamma):
|
|
if gamma == 1:
|
|
return tensor
|
|
if gamma == 2.2:
|
|
return torch.where(tensor <= .0031308,
|
|
tensor * 12.92,
|
|
1.055 * tensor ** (1/2.4) - 0.055)
|
|
return tensor ** (1/gamma)
|
|
|
|
def _grayscale_by_luminance(self, tensor, gamma):
|
|
weights = torch.tensor([0.2125, 0.7154, 0.0721], dtype=torch.float32)
|
|
tensor = self._gamma_expansion(tensor, gamma)
|
|
tensor = tensor @ weights
|
|
if gamma != 1:
|
|
tensor = tensor ** gamma
|
|
return self._gamma_compression(tensor, gamma)
|
|
|
|
def _squeeze_to_modules(self, tensor, method):
|
|
tensor = self._reshape_tensor_to_modules(tensor)
|
|
if method == "mean":
|
|
return self._squeeze_by_mean(tensor)
|
|
raise RuntimeError("Module aggregation currently only supports the "
|
|
"mean.")
|
|
|
|
def _reduce_to_modules(
|
|
self,
|
|
source_qr,
|
|
modified_qr,
|
|
module_size,
|
|
grayscale_method,
|
|
aggregate_method,
|
|
inverted_pattern,
|
|
gamma
|
|
):
|
|
self.module_size = module_size
|
|
self.canvas_shape = (source_qr.shape[0], source_qr.shape[1])
|
|
# Processed first for simplified indexing of QR bounds.
|
|
source_qr = self._squeeze_by_mean(source_qr)
|
|
source_qr, modified_qr = self._trim_to_qr_area(source_qr,
|
|
modified_qr,
|
|
inverted_pattern
|
|
)
|
|
if grayscale_method == "mean":
|
|
modified_qr = self._squeeze_by_mean(modified_qr)
|
|
elif grayscale_method == "luminance":
|
|
modified_qr = self._grayscale_by_luminance(modified_qr, gamma)
|
|
else:
|
|
raise ValueError("Currently only mean is supported for rgb to "
|
|
"grayscale conversion.")
|
|
source_qr = torch.round(self._squeeze_to_modules(source_qr, "mean"))
|
|
modified_qr = self._squeeze_to_modules(modified_qr, aggregate_method)
|
|
return source_qr, modified_qr
|
|
|
|
def _create_finder_pattern_mask(self, width, inverted):
|
|
mask = torch.zeros((width, width), dtype=torch.bool)
|
|
# When borders are trimmed and QR code has module size of 1, results
|
|
# are consistent.
|
|
finder_coords = [[0, 7, 0, 7], [0, 7, -7, None], [-7, None, 0, 7]]
|
|
for x_min, x_max, y_min, y_max in finder_coords:
|
|
mask[y_min:y_max, x_min:x_max] = True
|
|
return ~mask if inverted else mask
|
|
|
|
def _create_qr_mask(self, tensor, evaluate):
|
|
if evaluate == "module_pattern":
|
|
return self._create_finder_pattern_mask(tensor, True)
|
|
if evaluate == "finder_pattern":
|
|
return self._create_finder_pattern_mask(tensor, False)
|
|
return None
|
|
|
|
def _bin_tensor_to_threshold(self, tensor, contrast_difficulty):
|
|
tensor = tensor.clone()
|
|
threshold = contrast_difficulty / 2
|
|
# Since we are only interested in value matches and there is a clear
|
|
# stable dividing line of .5, bringing in the other array is
|
|
# unneccessary and the binning process can be simplified.
|
|
bin_condition = (tensor + threshold <= .5) & (tensor != .5)
|
|
tensor[bin_condition] = 0.0
|
|
bin_condition = (tensor - threshold >= .5) & (tensor != .5)
|
|
tensor[bin_condition] = 1.0
|
|
return tensor
|
|
|
|
def _replace_qr_to_canvas(self, tensor):
|
|
length = tensor.shape[0] * self.module_size
|
|
bounds = self.qr_bounds
|
|
tensor = F.interpolate(tensor.unsqueeze(0).unsqueeze(0),
|
|
size=(length, length),
|
|
mode='nearest')
|
|
canvas = torch.zeros(self.canvas_shape, dtype=torch.float32)
|
|
canvas[bounds[0]:bounds[1], bounds[2]:bounds[3]] = tensor.squeeze()
|
|
return canvas
|
|
|
|
def _compare_modules(
|
|
self,
|
|
source_qr,
|
|
modified_qr,
|
|
mask,
|
|
error_difficulty
|
|
):
|
|
modified_qr = self._bin_tensor_to_threshold(modified_qr,
|
|
error_difficulty)
|
|
error = source_qr != modified_qr
|
|
percent_error = error[mask].sum().item() / error[mask].numel()
|
|
if mask is not None:
|
|
error[~mask] = False
|
|
return (self._replace_qr_to_canvas((error).to(torch.float32)),
|
|
percent_error)
|
|
|
|
def _qr_correlation(self, source_qr, modified_qr, mask):
|
|
source_qr = source_qr[mask].numpy().reshape((-1))
|
|
modified_qr = modified_qr[mask].numpy().reshape((-1))
|
|
return np.corrcoef(source_qr, modified_qr)[0, 1]
|
|
|
|
def _qr_rmse(self, source_qr, modified_qr, mask):
|
|
diff = source_qr[mask].numpy() - modified_qr[mask].numpy()
|
|
return np.sqrt((diff ** 2).mean())
|
|
|
|
def find_qr_errors(
|
|
self,
|
|
source_qr,
|
|
modified_qr,
|
|
module_size,
|
|
grayscale_method,
|
|
aggregate_method,
|
|
evaluate,
|
|
error_difficulty,
|
|
inverted_pattern,
|
|
gamma,
|
|
):
|
|
self._check_equal_shape(source_qr, modified_qr)
|
|
error_masks, error_percents, correlations, rmses = [], [], [], []
|
|
for i in range(source_qr.shape[0]):
|
|
qr_s, qr_m = source_qr[i], modified_qr[i]
|
|
qr_s, qr_m = self._reduce_to_modules(qr_s,
|
|
qr_m,
|
|
module_size,
|
|
grayscale_method,
|
|
aggregate_method,
|
|
inverted_pattern,
|
|
gamma
|
|
)
|
|
mask = self._create_qr_mask(qr_s.shape[0], evaluate)
|
|
error_mask, percent_error = self._compare_modules(qr_s,
|
|
qr_m,
|
|
mask,
|
|
error_difficulty
|
|
)
|
|
correlation = self._qr_correlation(qr_s, qr_m, mask)
|
|
rmse = self._qr_rmse(qr_s, qr_m, mask)
|
|
error_masks.append(error_mask)
|
|
error_percents.append(percent_error)
|
|
correlations.append(correlation)
|
|
rmses.append(rmse)
|
|
error_masks = torch.stack(error_masks, dim=0)
|
|
return (error_masks, error_percents, correlations, rmses)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"comfy-qr-by-module-size": QRByModuleSize,
|
|
"comfy-qr-by-image-size": QRByImageSize,
|
|
"comfy-qr-by-module-split":
|
|
QRByModuleSizeSplitFunctionPatterns,
|
|
"comfy-qr-mask_errors": QRErrorMasker,
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"comfy-qr-by-module-size": "QR Code",
|
|
"comfy-qr-by-image-size": "QR Code (Conformed "
|
|
"to Image Size)",
|
|
"comfy-qr-by-module-split": "QR Code (Split)",
|
|
"comfy-qr-mask_errors": "Mask QR Errors",
|
|
}
|