148 lines
5.3 KiB
Python
148 lines
5.3 KiB
Python
from dataclasses import dataclass
|
|
import numpy as np
|
|
from PIL import Image
|
|
from skimage import transform, feature, measure
|
|
|
|
|
|
class ScaleTranslationTransform(transform.SimilarityTransform):
|
|
"""
|
|
A transformation that only considers uniform scaling and translation.
|
|
This is a specialized version of SimilarityTransform that ignores rotation and shear.
|
|
"""
|
|
|
|
def estimate(self, src, dst):
|
|
# src and dst are (N, 2) arrays
|
|
if src.shape[0] < 2:
|
|
return False
|
|
|
|
# 1. Compute centroids
|
|
src_mean = src.mean(axis=0)
|
|
dst_mean = dst.mean(axis=0)
|
|
|
|
# 2. Center the points
|
|
src_centered = src - src_mean
|
|
dst_centered = dst - dst_mean
|
|
|
|
# 3. Compute scale s
|
|
# s = sum(src_centered * dst_centered) / sum(src_centered^2)
|
|
# We sum over both dimensions and all points
|
|
num = np.sum(src_centered * dst_centered)
|
|
den = np.sum(src_centered**2)
|
|
|
|
if den == 0:
|
|
return False
|
|
|
|
s = num / den
|
|
|
|
# 4. Compute translation
|
|
# t = dst_mean - s * src_mean
|
|
t = dst_mean - s * src_mean
|
|
|
|
# 5. Set parameters
|
|
# Update the matrix
|
|
self.params = np.array([[s, 0, t[0]], [0, s, t[1]], [0, 0, 1]])
|
|
|
|
return True
|
|
|
|
|
|
@dataclass
|
|
class ImageWarper:
|
|
"""
|
|
Warp an edited image to align with an original image using feature matching.
|
|
|
|
This class uses ORB feature detection and RANSAC to find a robust
|
|
ScaleTranslationTransform. If alignment fails, it falls back to simple resizing.
|
|
"""
|
|
|
|
max_keypoints: int = 1000
|
|
min_matches: int = 10
|
|
ransac_residual_threshold: float = 2.0
|
|
|
|
def warp(self, original: Image.Image, edited: Image.Image) -> Image.Image:
|
|
"""
|
|
Aligns the edited image to the original image using feature matching and warping.
|
|
If alignment fails, falls back to simple resizing.
|
|
"""
|
|
# Convert to grayscale for feature detection
|
|
img1_gray = np.array(original.convert('L'))
|
|
img2_gray = np.array(edited.convert('L'))
|
|
|
|
# Detect ORB features
|
|
orb = feature.ORB(n_keypoints=self.max_keypoints)
|
|
|
|
try:
|
|
orb.detect_and_extract(img1_gray)
|
|
keypoints1 = orb.keypoints
|
|
descriptors1 = orb.descriptors
|
|
|
|
orb.detect_and_extract(img2_gray)
|
|
keypoints2 = orb.keypoints
|
|
descriptors2 = orb.descriptors
|
|
except RuntimeError:
|
|
return edited.resize(original.size, Image.Resampling.LANCZOS)
|
|
|
|
# Check if descriptors were found
|
|
if descriptors1 is None or descriptors2 is None or len(descriptors1) == 0 or len(descriptors2) == 0:
|
|
return edited.resize(original.size, Image.Resampling.LANCZOS)
|
|
|
|
# Match descriptors
|
|
matches12 = feature.match_descriptors(descriptors1, descriptors2, cross_check=True)
|
|
|
|
if len(matches12) < self.min_matches:
|
|
return edited.resize(original.size, Image.Resampling.LANCZOS)
|
|
|
|
# Extract location of good matches
|
|
# keypoints are (row, col)
|
|
src_rc = keypoints2[matches12[:, 1]] # edited
|
|
dst_rc = keypoints1[matches12[:, 0]] # original
|
|
|
|
# skimage transforms use (x, y) coordinates, which is (col, row).
|
|
# We need to swap the columns to work with the transform model.
|
|
src = src_rc[:, ::-1]
|
|
dst = dst_rc[:, ::-1]
|
|
|
|
# Estimate transform using RANSAC
|
|
# ScaleTranslationTransform handles translation and uniform scaling (no rotation).
|
|
model_robust, inliers = measure.ransac(
|
|
(src, dst), ScaleTranslationTransform, min_samples=3, residual_threshold=self.ransac_residual_threshold, max_trials=100
|
|
)
|
|
|
|
if model_robust is None:
|
|
#raise ValueError("RANSAC failed to find a model.")
|
|
return edited.resize(original.size, Image.Resampling.LANCZOS)
|
|
|
|
# print(f"Matches: {len(matches12)}")
|
|
# print(f"Inliers: {np.sum(inliers)}")
|
|
# print(f"Scale: {model_robust.scale:.4f}")
|
|
# print(f"Rotation: {model_robust.rotation:.4f} rad")
|
|
# print(f"Translation: {model_robust.translation}")
|
|
|
|
# Warp the edited image
|
|
edited_arr = np.array(edited.convert('RGBA'))
|
|
|
|
# To avoid black borders on the warped image, we'll warp RGB and Alpha channels separately.
|
|
# RGB channels are warped with 'edge' mode to prevent color bleeding from transparent areas.
|
|
# Alpha channel is warped with 'constant' mode (cval=0) to preserve transparency.
|
|
output_shape = np.array(original).shape[:2]
|
|
|
|
# Warp RGB channels
|
|
rgb_warped = transform.warp(edited_arr[..., :3], model_robust.inverse, output_shape=output_shape, order=1, mode='edge', preserve_range=True)
|
|
|
|
# Warp Alpha channel
|
|
alpha_warped = transform.warp(
|
|
edited_arr[..., 3], model_robust.inverse, output_shape=output_shape, order=1, mode='constant', cval=0, preserve_range=True
|
|
)
|
|
|
|
# Combine the warped RGB and alpha channels
|
|
warped_arr = np.dstack((rgb_warped, alpha_warped))
|
|
|
|
warped_img = Image.fromarray(warped_arr.astype(np.uint8), 'RGBA')
|
|
|
|
# Create a copy of the original image to paste onto
|
|
pasted_img = original.copy().convert('RGBA')
|
|
|
|
# Alpha composite the warped image onto the original
|
|
pasted_img = Image.alpha_composite(pasted_img, warped_img)
|
|
|
|
return pasted_img
|