custom_nodes/ComfyQR/qr_nodes.py

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",
                              }