christiangeorgelucas/image-hash-tools

from gen.messages_pb2 import AverageHashInput, HashResult
from gen.axiom_context import AxiomContext
from nodes._imaging import load_pil
from nodes._hashing import compute_square_hash, hash_result_fields


def average_hash(ax: AxiomContext, input: AverageHashInput) -> HashResult:
    """Average hash (aHash): resize to hash_size x hash_size grayscale, threshold
    each pixel against the mean. Fast and robust to scaling/gamma changes, but
    the least discriminative of the luminance hashes. hash_size defaults to 8
    (a 64-bit hash) when 0. Deterministic: the same image + hash_size always
    produce the same hex-encoded hash.
    """
    pil = load_pil(input.image)
    h, size = compute_square_hash("ahash", pil, input.hash_size)
    return HashResult(**hash_result_fields(h, "ahash", size))

from gen.messages_pb2 import ColorHashInput, HashResult
from gen.axiom_context import AxiomContext
from nodes._imaging import load_pil
from nodes._hashing import DEFAULT_BINBITS, hash_result_fields, resolve_size
import imagehash


def color_hash(ax: AxiomContext, input: ColorHashInput) -> HashResult:
    """Color hash: bins the black/gray fraction plus hue histograms of highly-
    and mildly-saturated pixels from the image's HSV histogram. Unlike the
    luminance hashes (aHash/dHash/pHash/wHash), this is position-independent —
    it hashes the color *distribution*, not a spatial grid, so it is useful
    for matching by color palette regardless of layout. binbits (bits of
    resolution per bin) defaults to 3 when 0. Deterministic: the same image +
    binbits always produce the same hex-encoded hash.
    """
    pil = load_pil(input.image)
    binbits = resolve_size(input.binbits, DEFAULT_BINBITS)
    h = imagehash.colorhash(pil, binbits=binbits)
    return HashResult(**hash_result_fields(h, "colorhash", binbits))

from gen.messages_pb2 import CropResistantHashInput, HashResult
from gen.axiom_context import AxiomContext
from nodes._imaging import load_pil
from nodes._hashing import (
    DEFAULT_MIN_SEGMENT_SIZE,
    DEFAULT_SEGMENTATION_IMAGE_SIZE,
    crop_resistant_hash_func,
    hash_result_fields,
    resolve_size,
)
import imagehash


def crop_resistant_hash(ax: AxiomContext, input: CropResistantHashInput) -> HashResult:
    """Crop-resistant hash: segments the image into bright/dark regions (a
    watershed-like split) and difference-hashes each segment, so the overall
    result stays matchable even after moderate cropping (the paper reports
    tolerance up to ~50% crop, vs ~5% for the single-hash algorithms above).
    The returned `hash` is a comma-joined list of per-segment hex hashes —
    reconstruct it with `imagehash.hex_to_multihash` if consuming outside this
    package. min_segment_size defaults to 500px, segmentation_image_size to
    300px, and the per-segment hash_size to 8, all when passed as 0.
    Deterministic: the same image + params always produce the same segments
    and hashes.
    """
    pil = load_pil(input.image)
    min_segment_size = resolve_size(input.min_segment_size, DEFAULT_MIN_SEGMENT_SIZE)
    segmentation_image_size = resolve_size(input.segmentation_image_size, DEFAULT_SEGMENTATION_IMAGE_SIZE)
    hash_func, size = crop_resistant_hash_func(input.hash_size)
    h = imagehash.crop_resistant_hash(
        pil,
        hash_func=hash_func,
        min_segment_size=min_segment_size,
        segmentation_image_size=segmentation_image_size,
    )
    return HashResult(**hash_result_fields(h, "crop_resistant", size))

from gen.messages_pb2 import DifferenceHashInput, HashResult
from gen.axiom_context import AxiomContext
from nodes._imaging import load_pil
from nodes._hashing import compute_square_hash, hash_result_fields


def difference_hash(ax: AxiomContext, input: DifferenceHashInput) -> HashResult:
    """Difference hash (dHash): resize to (hash_size+1) x hash_size grayscale,
    threshold each pixel against its left neighbor. Cheap and a good general-
    purpose near-duplicate hash. hash_size defaults to 8 (a 64-bit hash) when
    0. Deterministic: the same image + hash_size always produce the same
    hex-encoded hash.
    """
    pil = load_pil(input.image)
    h, size = compute_square_hash("dhash", pil, input.hash_size)
    return HashResult(**hash_result_fields(h, "dhash", size))

from gen.messages_pb2 import HashDistanceInput, HashDistanceResult
from gen.axiom_context import AxiomContext
from nodes._hashing import SQUARE_ALGORITHMS, reconstruct_colorhash, reconstruct_square_hash


def hash_distance(ax: AxiomContext, input: HashDistanceInput) -> HashDistanceResult:
    """Hamming distance between two hex-encoded hashes produced by this
    package's AverageHash/DifferenceHash/PerceptualHash/WaveletHash/ColorHash
    nodes. Lower distance means more similar; distance 0 means identical
    hashes. `algorithm` must match how BOTH hashes were computed ("ahash" |
    "dhash" | "phash" | "whash" | "colorhash"); for "colorhash", `hash_size`
    must equal the original `binbits`. crop_resistant hashes are not
    supported here — they need segment-aware comparison, not a flat hamming
    distance. Deterministic: the same two hashes always produce the same
    distance.
    """
    algorithm = input.algorithm
    if algorithm in SQUARE_ALGORITHMS:
        hash_a = reconstruct_square_hash(input.hash_a)
        hash_b = reconstruct_square_hash(input.hash_b)
    elif algorithm == "colorhash":
        hash_a = reconstruct_colorhash(input.hash_a, input.hash_size)
        hash_b = reconstruct_colorhash(input.hash_b, input.hash_size)
    else:
        raise ValueError(
            f"unknown algorithm {algorithm!r}; expected one of "
            f"{SQUARE_ALGORITHMS + ('colorhash',)}"
        )

    try:
        distance = hash_a - hash_b
    except TypeError as e:
        raise ValueError(f"hash_a and hash_b are not comparable: {e}") from e

    max_bits = len(hash_a)
    similarity = 1.0 - (distance / max_bits) if max_bits else 0.0
    return HashDistanceResult(distance=distance, max_bits=max_bits, similarity=similarity)

from gen.messages_pb2 import NearDuplicateInput, NearDuplicateResult
from gen.axiom_context import AxiomContext
from nodes._imaging import load_pil
from nodes._hashing import DEFAULT_ALGORITHM, DEFAULT_THRESHOLD, SQUARE_ALGORITHMS, compute_square_hash, resolve_size


def near_duplicate(ax: AxiomContext, input: NearDuplicateInput) -> NearDuplicateResult:
    """Convenience node: hashes both images with the same luminance algorithm
    ("ahash" | "dhash" | "phash" | "whash", default "phash" when empty) and
    hash_size (default 8), then reports whether their hamming distance is
    within `threshold` (default 10 — a commonly-used cutoff for a 64-bit
    luminance hash; lower is stricter). Equivalent to calling the matching
    hash node twice plus HashDistance, in one call. Deterministic: the same
    two images + params always produce the same verdict.
    """
    algorithm = input.algorithm or DEFAULT_ALGORITHM
    if algorithm not in SQUARE_ALGORITHMS:
        raise ValueError(f"unknown algorithm {algorithm!r}; expected one of {SQUARE_ALGORITHMS}")
    threshold = resolve_size(input.threshold, DEFAULT_THRESHOLD)

    pil_a = load_pil(input.image_a)
    pil_b = load_pil(input.image_b)
    hash_a, _ = compute_square_hash(algorithm, pil_a, input.hash_size)
    hash_b, _ = compute_square_hash(algorithm, pil_b, input.hash_size)
    distance = hash_a - hash_b

    return NearDuplicateResult(
        is_near_duplicate=distance <= threshold,
        distance=distance,
        threshold=threshold,
        algorithm=algorithm,
        hash_a=str(hash_a),
        hash_b=str(hash_b),
    )

from gen.messages_pb2 import PerceptualHashInput, HashResult
from gen.axiom_context import AxiomContext
from nodes._imaging import load_pil
from nodes._hashing import compute_square_hash, hash_result_fields


def perceptual_hash(ax: AxiomContext, input: PerceptualHashInput) -> HashResult:
    """Perceptual hash (pHash): 2D DCT of a downscaled grayscale image,
    threshold the low-frequency coefficients against their median. The most
    widely used general-purpose near-duplicate hash — more robust to color/
    contrast/compression changes than aHash or dHash. hash_size defaults to 8
    (a 64-bit hash) when 0. Deterministic: the same image + hash_size always
    produce the same hex-encoded hash.
    """
    pil = load_pil(input.image)
    h, size = compute_square_hash("phash", pil, input.hash_size)
    return HashResult(**hash_result_fields(h, "phash", size))

from gen.messages_pb2 import WaveletHashInput, HashResult
from gen.axiom_context import AxiomContext
from nodes._imaging import load_pil
from nodes._hashing import compute_square_hash, hash_result_fields


def wavelet_hash(ax: AxiomContext, input: WaveletHashInput) -> HashResult:
    """Wavelet hash (wHash): discrete wavelet transform (Haar or Daubechies-4)
    of a downscaled grayscale image, threshold the low-frequency coefficients
    against their median. Often more robust to small crops/rotations than
    pHash. hash_size must be a power of 2, defaults to 8 (a 64-bit hash) when
    0. `mode` is "haar" (default) or "db4". Deterministic: the same image +
    hash_size + mode always produce the same hex-encoded hash.
    """
    pil = load_pil(input.image)
    h, size = compute_square_hash("whash", pil, input.hash_size, mode=input.mode)
    return HashResult(**hash_result_fields(h, "whash", size))