autocatalog/data/preprocessing.py

import numpy as np
from PIL import Image

COLOR_IMAGE_SIZE = 128
COLOR_FEATURE_DIM = 37

def extract_color_features(image, image_size=COLOR_IMAGE_SIZE):
    image = image.convert("RGB").resize((image_size, image_size))
    margin = int(image_size * 0.10)
    image = image.crop(
        (
            margin,
            margin,
            image_size - margin,
            image_size - margin,
        )
    )

    rgb = np.asarray(image, dtype=np.float32) / 255.0
    hsv = np.asarray(image.convert("HSV"), dtype=np.float32) / 255.0

    rgb_flat = rgb.reshape(-1, 3)
    hsv_flat = hsv.reshape(-1, 3)

    saturation = hsv_flat[:, 1]
    value = hsv_flat[:, 2]
    foreground_mask = (saturation > 0.08) | (value < 0.92)

    if foreground_mask.sum() < 256:
        foreground_mask = np.ones(len(hsv_flat), dtype=bool)

    selected_rgb = rgb_flat[foreground_mask]
    selected_hsv = hsv_flat[foreground_mask]

    hue_hist, _ = np.histogram(
        selected_hsv[:, 0],
        bins=12,
        range=(0.0, 1.0),
    )

    saturation_hist, _ = np.histogram(
        selected_hsv[:, 1],
        bins=8,
        range=(0.0, 1.0),
    )

    value_hist, _ = np.histogram(
        selected_hsv[:, 2],
        bins=8,
        range=(0.0, 1.0),
    )

    hue_hist = hue_hist.astype(np.float32)
    saturation_hist = saturation_hist.astype(np.float32)
    value_hist = value_hist.astype(np.float32)

    hue_hist /= max(hue_hist.sum(), 1.0)
    saturation_hist /= max(saturation_hist.sum(), 1.0)
    value_hist /= max(value_hist.sum(), 1.0)

    rgb_mean = selected_rgb.mean(axis=0).astype(np.float32)
    rgb_std = selected_rgb.std(axis=0).astype(np.float32)
    rgb_median = np.median(selected_rgb, axis=0).astype(np.float32)

    features = np.concatenate(
        [
            hue_hist,
            saturation_hist,
            value_hist,
            rgb_mean,
            rgb_std,
            rgb_median,
        ]
    ).astype(np.float32)

    if features.shape[0] != COLOR_FEATURE_DIM:
        raise ValueError(
            f"Expected {COLOR_FEATURE_DIM} color features, "
            f"got {features.shape[0]}"
        )

    return features