Files
jonstreeterandClaude 36860b09dd Embed ColorMNet dependencies for self-contained installation
This fixes the ModuleNotFoundError issues that occurred during clean installations.

Changes:
- Embedded colormnet_source dependencies directly into repository as:
  - colormnet_model/ (model architecture)
  - colormnet_inference/ (inference pipeline)
  - colormnet_util/ (utility functions)
  - colormnet_dataset/ (data transforms and utilities)
- Updated all internal imports to use new embedded paths
- Removed external colormnet_source path dependency from __init__.py
- Added graceful fallback for spatial_correlation_sampler CUDA extension
- Fixed shape calculation bugs in PyTorch fallback implementation
- Added missing dependencies (hickle, scikit-image) to auto-install

The repository is now self-contained and can be installed without requiring
the separate colormnet_source directory.

Tested and verified:
- Node imports work correctly
- Colorization pipeline runs successfully
- Model loads and processes video frames

Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
2025-11-16 15:31:52 -06:00

38 lines
1.1 KiB
Python

import numpy as np
from PIL import Image
import cv2
import thinplate as tps
cv2.setNumThreads(0)
def pick_random_points(h, w, n_samples):
y_idx = np.random.choice(np.arange(h), size=n_samples, replace=False)
x_idx = np.random.choice(np.arange(w), size=n_samples, replace=False)
return y_idx/h, x_idx/w
def warp_dual_cv(img, mask, c_src, c_dst):
dshape = img.shape
theta = tps.tps_theta_from_points(c_src, c_dst, reduced=True)
grid = tps.tps_grid(theta, c_dst, dshape)
mapx, mapy = tps.tps_grid_to_remap(grid, img.shape)
return cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR), cv2.remap(mask, mapx, mapy, cv2.INTER_NEAREST)
def random_tps_warp(img, mask, scale, n_ctrl_pts=12):
"""
Apply a random TPS warp of the input image and mask
Uses randomness from numpy
"""
img = np.asarray(img)
mask = np.asarray(mask)
h, w = mask.shape
points = pick_random_points(h, w, n_ctrl_pts)
c_src = np.stack(points, 1)
c_dst = c_src + np.random.normal(scale=scale, size=c_src.shape)
warp_im, warp_gt = warp_dual_cv(img, mask, c_src, c_dst)
return Image.fromarray(warp_im), Image.fromarray(warp_gt)