rishicleaner / app.py
Miquel Farré
v1
3a79668
raw
history blame contribute delete
10.7 kB
import cv2
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import gradio as gr
import tempfile
import os
import shutil
def edge_directed_antialiasing(img, power=2.0):
"""
Apply edge-directed anti-aliasing with adjustable power
Parameters:
- img: Input image (numpy array)
- power: Anti-aliasing strength (1.0 is standard, higher values increase the effect)
Returns:
- Output image with anti-aliasing applied
"""
# If image has alpha channel, separate it
has_alpha = img.shape[2] == 4 if len(img.shape) > 2 else False
if has_alpha:
bgr = img[:, :, :3]
alpha = img[:, :, 3]
else:
bgr = img
# Create binary mask from grayscale image if no alpha
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
_, alpha = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
# Convert to grayscale for edge detection
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
# Step 1: Detect edges using Canny
# Lower thresholds to catch more edges when power is high
canny_threshold1 = int(100 / power) # Lower threshold when power is high
canny_threshold2 = int(200 / power) # Lower threshold when power is high
edges = cv2.Canny(gray, canny_threshold1, canny_threshold2)
# Dilate edges more when power is high
kernel_size = int(3 * power) # Increase kernel size with power
kernel_size = max(3, kernel_size if kernel_size % 2 == 1 else kernel_size + 1) # Ensure odd kernel size
kernel = np.ones((kernel_size, kernel_size), np.uint8)
# More iterations for higher power
dilation_iterations = max(1, int(power))
dilated_edges = cv2.dilate(edges, kernel, iterations=dilation_iterations)
# Step 2: Calculate gradient direction using Sobel
# Increase kernel size for higher power
sobel_ksize = 3
if power > 2.0:
sobel_ksize = 5
if power > 3.0:
sobel_ksize = 7
sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_ksize)
sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_ksize)
# Calculate gradient magnitude and direction
magnitude = np.sqrt(sobelx**2 + sobely**2)
direction = np.arctan2(sobely, sobelx) * 180 / np.pi
# Create output image, starting with the original
output = bgr.copy()
h, w = output.shape[:2]
# Step 3: Apply targeted smoothing along edge directions
# Sample farther away for higher power
radius = max(1, int(power))
edge_pixels = np.where(dilated_edges > 0)
for y, x in zip(edge_pixels[0], edge_pixels[1]):
# Skip border pixels
if x < radius or y < radius or x >= w-radius or y >= h-radius:
continue
# Get local direction (perpendicular to gradient)
local_dir = direction[y, x] + 90
if local_dir > 180:
local_dir -= 360
# Normalize direction to 0-180 degrees
local_dir = ((local_dir + 180) % 180)
# Determine interpolation direction based on edge angle
if 22.5 <= local_dir < 67.5: # ~45 degree diagonal
# Diagonal top-left to bottom-right
neighbors = [(y-radius, x-radius), (y+radius, x+radius)]
weights = [0.5, 0.5]
elif 67.5 <= local_dir < 112.5: # Vertical
# Top to bottom
neighbors = [(y-radius, x), (y+radius, x)]
weights = [0.5, 0.5]
elif 112.5 <= local_dir < 157.5: # ~135 degree diagonal
# Diagonal top-right to bottom-left
neighbors = [(y-radius, x+radius), (y+radius, x-radius)]
weights = [0.5, 0.5]
else: # Horizontal
# Left to right
neighbors = [(y, x-radius), (y, x+radius)]
weights = [0.5, 0.5]
# Only interpolate if we're between different colors (at the border)
center_value = gray[y, x]
neighbor_values = [gray[ny, nx] for ny, nx in neighbors]
# Lower contrast threshold when power is high
contrast_threshold = int(50 / power)
# Check if this is an edge between very different values
if abs(neighbor_values[0] - neighbor_values[1]) > contrast_threshold:
# Apply interpolation based on local contrast
for c in range(3): # RGB channels
weighted_sum = sum(weights[i] * bgr[ny, nx, c] for i, (ny, nx) in enumerate(neighbors))
# More interpolation weight when power is high
blend_factor = min(0.9, 0.3 * power)
# Apply it with a blend factor to preserve some original detail
output[y, x, c] = int((1-blend_factor) * weighted_sum + blend_factor * bgr[y, x, c])
# Update alpha channel with the same smoothing for edges
if has_alpha:
new_alpha = alpha.copy()
# Apply a specific smoothing to the alpha channel's edges
alpha_edges = cv2.Canny(alpha, int(100/power), int(200/power))
# More dilation iterations for stronger effect
alpha_dilation_iter = max(2, int(power * 2))
dilated_alpha_edges = cv2.dilate(alpha_edges, kernel, iterations=alpha_dilation_iter)
# Radius for sampling neighborhood
alpha_radius = max(2, int(power * 2))
# For each edge pixel in alpha
alpha_edge_pixels = np.where(dilated_alpha_edges > 0)
for y, x in zip(alpha_edge_pixels[0], alpha_edge_pixels[1]):
if x < alpha_radius or y < alpha_radius or x >= w-alpha_radius or y >= h-alpha_radius:
continue
# Use a larger neighborhood for better smoothing of alpha edges
# Size increases with power
window_radius = alpha_radius
neighborhood = alpha[y-window_radius:y+window_radius+1, x-window_radius:x+window_radius+1].astype(np.float32)
# Generate gaussian-like weights based on distance from center
kernel_size = 2 * window_radius + 1
weight_matrix = np.zeros((kernel_size, kernel_size), dtype=np.float32)
# Create distance-based weights
center = window_radius
for wy in range(kernel_size):
for wx in range(kernel_size):
# Calculate distance from center
dist = np.sqrt((wy - center)**2 + (wx - center)**2)
# Adjust falloff based on power
falloff = 1.0 / power
# Gaussian-like weight
weight_matrix[wy, wx] = np.exp(-(dist**2) / (2 * (window_radius * falloff)**2))
# Normalize weights
weight_matrix = weight_matrix / weight_matrix.sum()
# Apply weighted average
new_alpha[y, x] = int(np.sum(neighborhood * weight_matrix))
# Merge BGR with new alpha
output = np.dstack([output, new_alpha])
return output
def save_as_jpg(img, file_path):
"""
Save image as JPG with high quality
"""
# If image has alpha channel, blend with white background
if len(img.shape) > 2 and img.shape[2] == 4:
bgr = img[:, :, :3]
alpha = img[:, :, 3].astype(float) / 255
# Create white background
bg = np.ones_like(bgr) * 255
# Blend with background
alpha = np.expand_dims(alpha, axis=2)
alpha = np.repeat(alpha, 3, axis=2)
result = (bgr * alpha + bg * (1 - alpha)).astype(np.uint8)
else:
result = img
# Save as JPG
cv2.imwrite(file_path, result, [cv2.IMWRITE_JPEG_QUALITY, 95])
return file_path
def create_output_dirs():
"""Create necessary output directories"""
output_dir = os.path.join(tempfile.gettempdir(), "antialiasing_output")
os.makedirs(output_dir, exist_ok=True)
return output_dir
def process_image(input_image):
"""
Process image function for Gradio interface
"""
# Create output directory for our files
output_dir = create_output_dirs()
# Convert from RGB (Gradio) to BGR (OpenCV)
img_bgr = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
# Apply edge directed anti-aliasing with power=2.0
processed_bgr = edge_directed_antialiasing(img_bgr, power=2.0)
# Save the processed image explicitly as JPG
jpg_path = os.path.join(output_dir, "antialiased_image.jpg")
save_as_jpg(processed_bgr, jpg_path)
# Convert back to RGB for display in Gradio
if processed_bgr.shape[2] == 4: # Has alpha channel
# Blend with white background
bg = np.ones_like(processed_bgr[:,:,:3]) * 255
alpha = processed_bgr[:,:,3]
alpha_norm = alpha.astype(float) / 255
alpha_norm = np.expand_dims(alpha_norm, axis=2)
alpha_norm = np.repeat(alpha_norm, 3, axis=2)
processed_rgb = processed_bgr[:,:,:3] * alpha_norm + bg * (1 - alpha_norm)
processed_rgb = processed_rgb.astype(np.uint8)
else:
processed_rgb = cv2.cvtColor(processed_bgr, cv2.COLOR_BGR2RGB)
# Create comparison visualization
h, w = input_image.shape[:2]
dpi = 100
plt.figure(figsize=(w*2/dpi, h/dpi), dpi=dpi)
plt.subplot(1, 2, 1)
plt.imshow(input_image)
plt.title("Original")
plt.axis('off')
plt.subplot(1, 2, 2)
plt.imshow(processed_rgb)
plt.title("Anti-aliased (Power = 2.0)")
plt.axis('off')
plt.tight_layout()
# Save the comparison
comparison_file = os.path.join(output_dir, "comparison.jpg")
plt.savefig(comparison_file, dpi=dpi, bbox_inches='tight')
plt.close()
return processed_rgb, jpg_path, comparison_file
# Create Gradio interface
with gr.Blocks(title="Edge-Directed Anti-Aliasing") as app:
gr.Markdown("# Edge-Directed Anti-Aliasing Tool")
gr.Markdown("Upload an image and apply edge-directed anti-aliasing to smooth jagged edges.")
with gr.Row():
input_image = gr.Image(label="Upload Image", type="numpy")
output_image = gr.Image(label="Anti-Aliased Result", type="numpy")
with gr.Row():
process_button = gr.Button("Apply Anti-Aliasing (Power = 2.0)")
with gr.Row():
download_jpg = gr.File(label="Download Anti-Aliased JPG", type="filepath")
comparison_view = gr.Image(label="Comparison", type="filepath")
# Process button functionality
process_button.click(
fn=process_image,
inputs=[input_image],
outputs=[output_image, download_jpg, comparison_view]
)
# Launch the app
if __name__ == "__main__":
app.launch(share=True)