import numpy as np import traceback import gradio as gr from PIL import Image from scipy import ndimage, interpolate import matplotlib.pyplot as plt from bulk_bulge_generation import definitions, smooth # from transformers import pipeline import fastai from fastcore.all import * from fastai.vision.all import * from ultralytics import YOLO def apply_vector_field_transform(image, func, radius, center=(0.5, 0.5), strength=1, edge_smoothness=0.1, center_smoothness=0.20): rows, cols = image.shape[:2] max_dim = max(rows, cols) center_y = int(center[1] * rows) center_x = int(center[0] * cols) center_y = abs(rows - center_y) print(f"Image shape: {rows}x{cols}") print(f"Center: ({center_x}, {center_y})") print(f"Radius: {radius}, Strength: {strength}") print(f"Edge smoothness: {edge_smoothness}, Center smoothness: {center_smoothness}") y, x = np.ogrid[:rows, :cols] y = (y - center_y) / max_dim x = (x - center_x) / max_dim dist_from_center = np.sqrt(x**2 + y**2) z = func(x, y) print(f"Function output - min: {np.min(z)}, max: {np.max(z)}") gy, gx = np.gradient(z) print(f"Initial gradient - gx min: {np.min(gx)}, max: {np.max(gx)}") print(f"Initial gradient - gy min: {np.min(gy)}, max: {np.max(gy)}") # Avoid division by zero edge_smoothness = np.maximum(edge_smoothness, 1e-6) center_smoothness = np.maximum(center_smoothness, 1e-6) edge_mask = np.clip((radius - dist_from_center) / (radius * edge_smoothness), 0, 1) center_mask = np.clip((dist_from_center - radius * center_smoothness) / (radius * center_smoothness), 0, 1) mask = edge_mask * center_mask gx = gx * mask gy = gy * mask magnitude = np.sqrt(gx**2 + gy**2) magnitude[magnitude == 0] = 1 # Avoid division by zero gx = gx / magnitude gy = gy / magnitude scale_factor = strength * np.log(max_dim) / 100 gx = gx * scale_factor * mask gy = gy * scale_factor * mask print(f"Final gradient - gx min: {np.min(gx)}, max: {np.max(gx)}") print(f"Final gradient - gy min: {np.min(gy)}, max: {np.max(gy)}") # Forward transformation x_new = x + gx y_new = y + gy x_new = x_new * max_dim + center_x y_new = y_new * max_dim + center_y x_new = np.clip(x_new, 0, cols - 1) y_new = np.clip(y_new, 0, rows - 1) # Inverse transformation x_inv = x - gx y_inv = y - gy x_inv = x_inv * max_dim + center_x y_inv = y_inv * max_dim + center_y x_inv = np.clip(x_inv, 0, cols - 1) y_inv = np.clip(y_inv, 0, rows - 1) # Apply transformations channels_forward = [ndimage.map_coordinates(image[..., i], [y_new, x_new], order=1, mode='reflect') for i in range(image.shape[2])] channels_inverse = [ndimage.map_coordinates(image[..., i], [y_inv, x_inv], order=1, mode='reflect') for i in range(image.shape[2])] transformed_image = np.dstack(channels_forward).astype(image.dtype) inverse_transformed_image = np.dstack(channels_inverse).astype(image.dtype) return transformed_image, inverse_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 import numpy as np from scipy import interpolate # def invert_gradient_vector_field(gx, gy, image_shape): # """ # Invert the gradient vector field using a more accurate method. # # :param gx: X-component of the gradient # :param gy: Y-component of the gradient # :param image_shape: Shape of the original image (height, width) # :return: Inverted gx and gy # """ # rows, cols = image_shape # y, x = np.mgrid[0:rows, 0:cols] # # # Calculate the new positions after applying the gradient # new_x = x + gx # new_y = y + gy # # # Create a mask for valid (non-NaN, non-infinite) values # mask = np.isfinite(new_x) & np.isfinite(new_y) # # # Flatten and filter the arrays # x_flat = x[mask] # y_flat = y[mask] # new_x_flat = new_x[mask] # new_y_flat = new_y[mask] # # # Create the inverse mapping # inv_x = interpolate.griddata((new_x_flat, new_y_flat), x_flat, (x, y), method='linear', fill_value=np.nan) # inv_y = interpolate.griddata((new_x_flat, new_y_flat), y_flat, (x, y), method='linear', fill_value=np.nan) # # # Calculate the inverse gradient # inv_gx = inv_x - x # inv_gy = inv_y - y # # # Fill NaN values with zeros # inv_gx = np.nan_to_num(inv_gx) # inv_gy = np.nan_to_num(inv_gy) # # return -inv_gx, -inv_gy # Note the negation here def apply_gradient_transform(image, gx, gy): """ Apply the gradient transformation to an image. :param image: Input image as a numpy array :param gx: X-component of the gradient :param gy: Y-component of the gradient :return: Transformed image """ rows, cols = image.shape[:2] y, x = np.mgrid[0:rows, 0:cols] # Apply the transformation x_new = x + gx y_new = y + gy # 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 = [] for i in range(image.shape[2]): channel = image[:,:,i] transformed_channel = interpolate.griddata((y.flatten(), x.flatten()), channel.flatten(), (y_new, x_new), method='linear', fill_value=0) channels.append(transformed_channel) transformed_image = np.dstack(channels).astype(image.dtype) return transformed_image ############################# # MAIN FUNCTION HERE ############################# # Version Check print(f"NumPy version: {np.__version__}") print(f"PyTorch version: {torch.__version__}") print(f"FastAI version: {fastai.__version__}") learn_bias = load_learner('model_bias.pkl') learn_fresh = load_learner('model_fresh.pkl') # Loads the YOLO Model model = YOLO("bulge_yolo_model.pt") def transform_image(image, func_choice, randomization_check, radius, center_x, center_y, strength, 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 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) rng = np.random.default_rng() if randomization_check: radius, location, strength, edge_smoothness = definitions(rng) center_x, center_y = location else: edge_smoothness, center_smoothness = smooth(rng, strength) if func_choice == "Pinch": func = pinch elif func_choice == "Spiral": func = shift elif func_choice == "Bulge": func = bulge edge_smoothness = 0 center_smoothness = 0 elif func_choice == "Volcano": func = bulge elif func_choice == "Shift Up": func = lambda x, y: spiral(x, y, frequency=spiral_frequency) print(f"Function choice: {func_choice}") print(f"Input image shape: {I.shape}") try: transformed, inverse_transformed, (gx, gy) = apply_vector_field_transform( I, func, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness ) print(f"Transformed image shape: {transformed.shape}") print(f"Inverse transformed image shape: {inverse_transformed.shape}") print(f"Gradient shapes: gx {gx.shape}, gy {gy.shape}") print(f"Gradient ranges: gx [{np.min(gx)}, {np.max(gx)}], gy [{np.min(gy)}, {np.max(gy)}]") vector_field = create_gradient_vector_field(gx, gy, I.shape[:2], reverse=reverse_gradient) inverted_vector_field = create_gradient_vector_field(-gx, -gy, I.shape[:2], reverse=False) print(f"Vector field shape: {vector_field.shape}") print(f"Inverted vector field shape: {inverted_vector_field.shape}") except Exception as e: print(f"Error in transformation: {str(e)}") traceback.print_exc() transformed = np.zeros_like(I) inverse_transformed = np.zeros_like(I) vector_field = np.zeros_like(I) inverted_vector_field = np.zeros_like(I) result = Image.fromarray(transformed) categories = ['Distorted', 'Maze'] def clean_output(result_values): pred, idx, probs = result_values return dict(zip(categories, map(float, probs))) result_bias = learn_bias.predict(result) result_fresh = learn_fresh.predict(result) result_bias_final = clean_output(result_bias) result_fresh_final = clean_output(result_fresh) result_localization = model.predict(transformed, save=True) return transformed, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field demo = gr.Interface( fn=transform_image, inputs=[ gr.Image(type="filepath"), gr.Dropdown(["Pinch", "Spiral", "Shift Up", "Bulge", "Volcano"], value="Volcano", label="Function"), gr.Checkbox(label="Randomize inputs?"), 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.Slider(0, 0.5, value=0.1, label="Center Smoothness") # gr.Checkbox(label="Reverse Gradient Direction"), ], examples=[ [np.asarray(Image.open("examples/1500_maze.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5], [np.asarray(Image.open("examples/2048_maze.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5], [np.asarray(Image.open("examples/2300_fresh.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5], [np.asarray(Image.open("examples/50_fresh.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5] ], outputs=[ gr.Image(label="Transformed Image"), gr.Label(), gr.Label(), gr.Image(label="Gradient Vector Field"), gr.Image(label="Inverse Gradient"), gr.Image(label="Inverted Vector Field"), ], title="Image Transformation Demo!", article="If you like this demo, please star the github repository for the project! Located [here!](https://github.com/nick-leland/DistortionML)", 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!\nYou 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'/'Volcano' function, I can't really guarantee that the others work as well!\nI have just added the first baseline ML model to detect if a distortion has taken place! It was only trained on mazes though ([Dataset Here](https://www.kaggle.com/datasets/nickleland/distorted-mazes)) so in order for it to detect a distortion you have to use one of the images provided in the examples! Feel free to mess around wtih other images in the meantime though!" ) demo.launch(share=True)