control_maps.py

import random

import cv2
import numpy as np
import torch


def nms(x, threshold=127, sigma=3.0):
    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), sigma)
    kernels = [
        np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8),
        np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8),
        np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8),
        np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8),
    ]
    y = np.zeros_like(x)
    for kernel in kernels:
        np.putmask(y, cv2.dilate(x, kernel=kernel) == x, x)
    out = np.zeros_like(y, dtype=np.uint8)
    out[y > threshold] = 255
    return out


def hed_to_scribble(edge_u8, thickness=None):
    """Convert soft HED output to Sana-style sparse scribble control."""
    if edge_u8.ndim == 3:
        edge_u8 = cv2.cvtColor(edge_u8, cv2.COLOR_RGB2GRAY)
    edge = nms(edge_u8, threshold=127, sigma=3.0)
    edge = cv2.GaussianBlur(edge, (0, 0), 3.0)
    edge[edge > 4] = 255
    edge[edge < 255] = 0

    if thickness is None:
        thickness = random.randint(0, 10)
    if thickness == 0:
        kernel = np.ones((3, 3), np.uint8)
        edge = cv2.erode(edge, kernel, iterations=1)
    elif thickness > 1:
        kernel_size = max(1, thickness // 2)
        kernel = np.ones((kernel_size, kernel_size), np.uint8)
        edge = cv2.dilate(edge, kernel, iterations=1)
    return edge.astype(np.uint8)


def control_to_tensor(edge_u8):
    if edge_u8.ndim == 2:
        edge_u8 = np.stack([edge_u8, edge_u8, edge_u8], axis=0)
    else:
        edge_u8 = edge_u8.transpose(2, 0, 1)
    arr = edge_u8.astype(np.float32) / 127.5 - 1.0
    return torch.from_numpy(arr)