Spaces:
Running on Zero
Running on Zero
| """Optional CUDA for the two heavy numeric steps (colour clustering, per-label smoothing), and the Hugging Face ZeroGPU hook. | |
| Nothing here is required: with no torch, no CUDA device, or HATCHLING_GPU=0, ling_hatch uses scikit-learn and OpenCV as before. | |
| The GPU path gives the same kind of result, not bit-identical labels, so a sheet made on the GPU can differ in a few edge pixels | |
| from the same sheet made on the CPU. | |
| """ | |
| import functools | |
| import math | |
| import os | |
| try: | |
| import spaces # present on Hugging Face Spaces; on ZeroGPU the GPU is only attached inside @spaces.GPU calls | |
| except Exception: | |
| spaces = None | |
| try: | |
| import torch | |
| import torch.nn.functional as F | |
| except Exception: | |
| torch = None | |
| MIN_PIXELS = 20000 # smaller label maps are faster on the CPU than the trip to the device | |
| _active = 0 # how many gpu_task calls are running in this process | |
| _failed = False # CUDA raised once: stay on the CPU from then on | |
| # On a Hugging Face Space with ZeroGPU, torch pretends to have CUDA at all times but only really has it inside a @spaces.GPU call: outside one, the | |
| # first CUDA operation ends in "Low-level CUDA init reached". So on a Space the GPU is used only inside a gpu_task. | |
| ON_SPACE = spaces is not None and bool(os.environ.get('SPACE_ID') or os.environ.get('SPACES_ZERO_GPU')) | |
| def device(): | |
| """'cuda' when CUDA is usable and not switched off, else None. On a Space this is only so inside a gpu_task call.""" | |
| if torch is None or _failed or os.environ.get('HATCHLING_GPU', '1').lower() in ('0', 'false', 'no', 'off'): | |
| return None | |
| if ON_SPACE and not _active: return None | |
| try: | |
| return 'cuda' if torch.cuda.is_available() else None | |
| except Exception: | |
| return None | |
| def fail(e): | |
| """CUDA refused: stay on the CPU from here on.""" | |
| global _failed | |
| _failed = True; print('GPU path failed (%s: %s): using the CPU' % (type(e).__name__, str(e)[:120])) | |
| def attempt(on_gpu, on_cpu): | |
| """Run the GPU version; if CUDA refuses for any reason, remember that and run the CPU version instead.""" | |
| if device(): | |
| try: return on_gpu() | |
| except Exception as e: fail(e) | |
| return on_cpu() | |
| def gpu_task(duration=120): | |
| """Decorator for the functions that reach the heavy steps: ZeroGPU's @spaces.GPU where it exists, otherwise nothing. It marks the call as one in | |
| which the GPU may be used.""" | |
| def wrap(fn): | |
| def marked(*a, **k): | |
| global _active | |
| _active += 1 | |
| try: return fn(*a, **k) | |
| finally: _active -= 1 | |
| return spaces.GPU(duration=duration)(marked) if spaces is not None else marked | |
| return wrap | |
| class KMeans: | |
| """k-means++ with n_init restarts on the GPU; the part of sklearn's KMeans that ling_hatch uses.""" | |
| def __init__(self, k, n_init=4, random_state=0, iters=60): | |
| self.k, self.n_init, self.seed, self.iters = k, n_init, int(random_state), iters | |
| self.cluster_centers_ = None | |
| def fit(self, X): | |
| dev = device() | |
| x = torch.as_tensor(X, dtype=torch.float32, device=dev) | |
| best = None | |
| for r in range(self.n_init): | |
| g = torch.Generator(device=dev).manual_seed(self.seed * 1009 + r) | |
| c = self._init(x, g) | |
| for _ in range(self.iters): | |
| a = torch.cdist(x, c).argmin(1) | |
| nc = torch.zeros_like(c).index_add_(0, a, x) | |
| cnt = torch.bincount(a, minlength=self.k).clamp_min(1).unsqueeze(1) | |
| nc = torch.where(torch.bincount(a, minlength=self.k).unsqueeze(1) > 0, nc / cnt, c) | |
| moved = (nc - c).abs().max().item(); c = nc | |
| if moved < 1e-4: break | |
| inertia = torch.cdist(x, c).min(1).values.pow(2).sum().item() | |
| if best is None or inertia < best[0]: best = (inertia, c) | |
| self._c = best[1]; self.cluster_centers_ = best[1].cpu().numpy().astype('float64') | |
| return self | |
| def _init(self, x, g): | |
| n = x.shape[0] | |
| c = x[torch.randint(n, (1,), generator=g, device=x.device)] | |
| for _ in range(1, self.k): | |
| d2 = torch.cdist(x, c).min(1).values.pow(2) | |
| tot = d2.sum() | |
| i = torch.multinomial(d2 / tot, 1, generator=g) if tot > 0 else torch.randint(n, (1,), generator=g, device=x.device) | |
| c = torch.cat([c, x[i]]) | |
| return c | |
| def predict(self, X): | |
| x = torch.as_tensor(X, dtype=torch.float32, device=self._c.device) | |
| out = [torch.cdist(x[i:i + (1 << 20)], self._c).argmin(1) for i in range(0, x.shape[0], 1 << 20)] | |
| return torch.cat(out).cpu().numpy().astype('int32') | |
| def _kernel(sigma, dev): | |
| """OpenCV's GaussianBlur kernel for float32 with ksize 0: round(8 sigma + 1) | 1 taps.""" | |
| n = int(round(sigma * 8 + 1)) | 1; r = n // 2 | |
| t = torch.arange(-r, r + 1, device=dev, dtype=torch.float32) | |
| k = torch.exp(-t * t / (2 * sigma * sigma)); return k / k.sum(), r | |
| def _blur(x, sigma): | |
| """x: (1, 1, H, W) float32. Separable Gaussian with reflect-101 borders, as cv2.GaussianBlur.""" | |
| k, r = _kernel(sigma, x.device) | |
| r = min(r, x.shape[2] - 1, x.shape[3] - 1) | |
| if r < 1: return x | |
| k = k[len(k) // 2 - r: len(k) // 2 + r + 1]; k = k / k.sum() | |
| x = F.conv2d(F.pad(x, (r, r, 0, 0), mode='reflect'), k.view(1, 1, 1, -1)) | |
| return F.conv2d(F.pad(x, (0, 0, r, r), mode='reflect'), k.view(1, 1, -1, 1)) | |
| def mode_clean(lbl, sigma, size=None): | |
| """Each pixel takes the label with the strongest Gaussian-smoothed presence (cv2 loop of _mode_clean / the upscale in posterize). | |
| size=(w, h): the result is bicubically upscaled to that size first, as posterize does for the working mask.""" | |
| import numpy as np | |
| dev = device() | |
| t = torch.as_tensor(lbl, device=dev) | |
| best = torch.full(size[::-1] if size else tuple(t.shape), -1.0 if size else 0.0, device=dev) | |
| out = torch.zeros(tuple(best.shape), dtype=torch.uint8, device=dev) if size else t.clone() | |
| for v in torch.unique(t).tolist(): | |
| sc = _blur((t == v).float()[None, None], sigma) | |
| if size: sc = F.interpolate(sc, size=(size[1], size[0]), mode='bicubic', align_corners=False) | |
| sc = sc[0, 0]; m = sc > best | |
| out[m] = v; best = torch.where(m, sc, best) | |
| return out.cpu().numpy().astype(np.uint8) | |