Spaces:
Running
Running
| 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) |