app.py

import numpy as np
import gradio as gr
from PIL import Image
from scipy import ndimage
import matplotlib.pyplot as plt

def apply_vector_field_transform(image, func, radius, center=(0.5, 0.5), strength=1, edge_smoothness=0.1):
    """
    Apply a vector field transformation to an image based on a given multivariate function.
    
    :param image: Input image as a numpy array (height, width, channels)
    :param func: A function that takes x and y as inputs and returns a scalar
    :param radius: Radius of the effect as a fraction of the image size
    :param center: Tuple (y, x) for the center of the effect, normalized to [0, 1]
    :param strength: Strength of the effect, scaled to image size
    :param edge_smoothness: Width of the smooth transition at the edge, as a fraction of the radius
    :return: Tuple of (transformed image as a numpy array, gradient vectors for vector field)
    """
    rows, cols = image.shape[:2]
    max_dim = max(rows, cols)
    
    # Convert normalized center to pixel coordinates
    center_y = int(center[0] * rows)
    center_x = int(center[1] * cols)
    
    # Convert normalized radius to pixel radius
    pixel_radius = int(max_dim * radius)
    
    y, x = np.ogrid[:rows, :cols]
    y = (y - center_y) / max_dim
    x = (x - center_x) / max_dim
    
    # Calculate distance from center
    dist_from_center = np.sqrt(x**2 + y**2)
    
    # Calculate function values
    z = func(x, y)
    
    # Calculate gradients
    gy, gx = np.gradient(z)

    # Create smooth transition mask
    mask = np.clip((radius - dist_from_center) / (radius * edge_smoothness), 0, 1)
    
    # Apply mask to gradients
    gx = gx * mask
    gy = gy * mask
    
    # Normalize gradient vectors
    magnitude = np.sqrt(gx**2 + gy**2)
    magnitude[magnitude == 0] = 1  # Avoid division by zero
    gx = gx / magnitude
    gy = gy / magnitude
    
    # Scale the effect (Play with the number 5)
    scale_factor = strength * np.log(max_dim) / 100  # Adjust strength based on image size
    gx = gx * scale_factor * mask
    gy = gy * scale_factor * mask
    
    # Create the mapping
    x_new = x + gx
    y_new = y + gy
    
    # Convert back to pixel coordinates
    x_new = x_new * max_dim + center_x
    y_new = y_new * max_dim + center_y
    
    # Ensure the new coordinates are within the image boundaries
    x_new = np.clip(x_new, 0, cols - 1)
    y_new = np.clip(y_new, 0, rows - 1)
    
    # Apply the transformation to each channel
    channels = [ndimage.map_coordinates(image[..., i], [y_new, x_new], order=1, mode='reflect') 
                for i in range(image.shape[2])]
    
    transformed_image = np.dstack(channels).astype(image.dtype)
    
    return transformed_image, (gx, gy)

def create_gradient_vector_field(gx, gy, image_shape, step=20, reverse=False):
    """
    Create a gradient vector field visualization with option to reverse direction.
    
    :param gx: X-component of the gradient
    :param gy: Y-component of the gradient
    :param image_shape: Shape of the original image (height, width)
    :param step: Spacing between arrows
    :param reverse: If True, reverse the direction of the arrows
    :return: Gradient vector field as a numpy array (RGB image)
    """
    rows, cols = image_shape
    y, x = np.mgrid[step/2:rows:step, step/2:cols:step].reshape(2, -1).astype(int)
    
    # Calculate the scale based on image size
    max_dim = max(rows, cols)
    scale = max_dim / 1000  # Adjusted for longer arrows
    
    # Reverse direction if specified
    direction = -1 if reverse else 1
    
    fig, ax = plt.subplots(figsize=(cols/50, rows/50), dpi=100)
    ax.quiver(x, y, direction * gx[y, x], direction * -gy[y, x], 
              scale=scale, 
              scale_units='width', 
              width=0.002 * max_dim / 500,
              headwidth=8, 
              headlength=12, 
              headaxislength=0, 
              color='black',
              minshaft=2,
              minlength=0,
              pivot='tail')
    ax.set_xlim(0, cols)
    ax.set_ylim(rows, 0)
    ax.set_aspect('equal')
    ax.axis('off')
    
    fig.tight_layout(pad=0)
    fig.canvas.draw()
    vector_field = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    vector_field = vector_field.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    plt.close(fig)
    
    return vector_field

def transform_image(image, func_choice, radius, center_x, center_y, strength, edge_smoothness, reverse_gradient=True, spiral_frequency=1):
    I = np.asarray(Image.open(image))    

    def pinch(x, y):
        return x**2 + y**2

    def shift(x, y):
        return np.arctan2(y, x)

    def bulge(x, y):
        r = np.sqrt(x**2 + y**2)
        # return -1 / (r + 1)
        return -r 

    def spiral(x, y, frequency=1):
        r = np.sqrt(x**2 + y**2)
        theta = np.arctan2(y, x)
        return r * np.sin(theta - frequency * r)

    if func_choice == "Pinch":
        func = pinch
    elif func_choice == "Spiral":
        func = shift 
    elif func_choice == "Bulge":
        func = bulge
    elif func_choice == "Shift":
        func = lambda x, y: spiral(x, y, frequency=spiral_frequency)

    transformed, (gx, gy) = apply_vector_field_transform(I, func, radius, (center_y, center_x), strength, edge_smoothness)
    vector_field = create_gradient_vector_field(gx, gy, I.shape[:2], reverse=reverse_gradient)

    return transformed, vector_field

demo = gr.Interface(
    fn=transform_image,
    inputs=[
        gr.Image(type="filepath"),
        gr.Dropdown(["Pinch", "Spiral", "Shift", "Bulge"], value="Bulge", label="Function"), 
        gr.Slider(0, 0.5, value=0.25, label="Radius (as fraction of image size)"),
        gr.Slider(0, 1, value=0.5, label="Center X"),
        gr.Slider(0, 1, value=0.5, label="Center Y"),
        gr.Slider(0, 1, value=0.5, label="Strength"),
        gr.Slider(0, 1, value=0.5, label="Edge Smoothness")
        # gr.Checkbox(label="Reverse Gradient Direction"),
    ],
    outputs=[
        gr.Image(label="Transformed Image"),
        gr.Image(label="Gradient Vector Field")
    ],
    title="Image Transformation Demo!",
    description="This is the baseline function that will be used to generate the database for a machine learning model I am working on called 'DistortionMl'! The goal of this model is to detect and then reverse image transformations that can be generated here! You can read more about the project at this repository link : https://github.com/nick-leland/DistortionML. The main function that I was working on is the 'Bulge' function, I can't really guarantee that the others work well (;"
)

demo.launch(share=True)