98 lines
3.7 KiB
Python
98 lines
3.7 KiB
Python
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
from sklearn.cluster import MiniBatchKMeans
|
|
from sklearn.utils import shuffle
|
|
|
|
|
|
from .ld_processor import calc_ciede
|
|
from .ld_processor_np import get_cls_update
|
|
|
|
|
|
def get_blur_torch(img: torch.Tensor, labels: torch.Tensor, size, blur=True):
|
|
if blur:
|
|
assert size % 2 == 1
|
|
p = (size - 1) // 2
|
|
img = F.pad(img, [p, p, p, p], mode='reflect')
|
|
img = F.avg_pool2d(img, kernel_size=size, stride=1)
|
|
|
|
cls = torch.unique(labels).reshape(-1, 1, 1, 1)
|
|
masks = torch.bitwise_and(img[:, [3]] > 127, cls == labels)
|
|
|
|
cls_counts = masks.sum(dim=(2, 3), keepdim=True) + 1e-7
|
|
rgb_means = (img[:, :3] * masks).sum(dim=(2, 3), keepdim=True) / cls_counts
|
|
|
|
rgb_means = rgb_means.squeeze().cpu().tolist()
|
|
cls_list = cls.squeeze().cpu().tolist()
|
|
cls_counts = cls_counts.squeeze().cpu().tolist()
|
|
|
|
return rgb_means, cls_list, cls_counts, masks
|
|
|
|
|
|
def get_base_torch(img: np.ndarray, loop, cls_num, threshold, size, kmeans_samples=-1, device='cpu'):
|
|
rgb_flatten = cluster_samples = img[..., :3].reshape((-1, 3))
|
|
im_h, im_w = img.shape[:2]
|
|
|
|
alpha_mask = np.where(img[..., 3] > 127)
|
|
resampled = False
|
|
if rgb_flatten.shape[0] > len(alpha_mask[0]):
|
|
cluster_samples = img[..., :3][alpha_mask].reshape((-1, 3))
|
|
resampled = True
|
|
|
|
if len(rgb_flatten) > kmeans_samples and kmeans_samples > 0:
|
|
cluster_samples = shuffle(cluster_samples, random_state=0, n_samples=kmeans_samples)
|
|
resampled = True
|
|
|
|
kmeans = MiniBatchKMeans(n_clusters=cls_num).fit(cluster_samples)
|
|
if resampled:
|
|
labels = kmeans.predict(rgb_flatten)
|
|
else:
|
|
labels = kmeans.labels_
|
|
|
|
img_torch = rearrange([img], 'n h w c -> n c h w')
|
|
img_torch = torch.from_numpy(img_torch).to(dtype=torch.float32, device=device)
|
|
labels_torch = torch.from_numpy(labels.reshape((1, 1, im_h, im_w))).to(dtype=torch.float32, device=device)
|
|
|
|
assert loop > 0
|
|
img_torch_ori = img_torch.clone()
|
|
for i in range(loop):
|
|
rgb_means, cls_list, cls_counts, masks = get_blur_torch(img_torch, labels_torch, size)
|
|
ciede_df = calc_ciede(rgb_means, cls_list)
|
|
cls2rgb, cls2counts, cls2masks = {}, {}, {}
|
|
for c, rgb, count, mask in zip(cls_list, rgb_means, cls_counts, masks):
|
|
cls2rgb[c] = rgb
|
|
cls2counts[c] = count
|
|
cls2masks[c] = mask[None, ...]
|
|
merge_dict = get_cls_update(ciede_df, threshold, cls2counts)
|
|
tgt2merge, notmerged = {}, set(cls_list)
|
|
for k, v in merge_dict.items():
|
|
if v not in tgt2merge:
|
|
tgt2merge[v] = []
|
|
notmerged.remove(v)
|
|
tgt2merge[v].append(k)
|
|
notmerged.remove(k)
|
|
for k in notmerged:
|
|
tgt2merge[k] = []
|
|
|
|
for tgtc, srcc_list in tgt2merge.items():
|
|
mask = cls2masks[tgtc]
|
|
for srcc in srcc_list:
|
|
mask = torch.bitwise_or(mask, cls2masks[srcc])
|
|
labels_torch.masked_fill_(mask, tgtc)
|
|
if i != loop - 1:
|
|
for jj in range(3):
|
|
img_torch[:, jj].masked_fill_(mask[0], cls2rgb[tgtc][jj])
|
|
|
|
cls_list = torch.unique(labels_torch)
|
|
img_torch = img_torch_ori
|
|
rgb_means, cls_list, cls_counts, masks = get_blur_torch(img_torch, labels_torch, size, blur=False)
|
|
for mask, rgb in zip(masks, rgb_means):
|
|
for jj in range(3):
|
|
img_torch[:, jj][mask] = rgb[jj]
|
|
|
|
img = rearrange(img_torch.cpu().numpy(), 'n c h w -> h w (n c)')
|
|
img = img.clip(0, 255).astype(np.uint8)
|
|
labels = labels_torch.cpu().numpy().squeeze().astype(np.uint32)
|
|
return img, labels |