Updated the gradio app to now include the inverse gradient generation

app.py

@@ -1,100 +1,95 @@
 import numpy as np
+import traceback
 import gradio as gr
 from PIL import Image
-from scipy import ndimage
+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):
-    # 0.106 strength = .50
-    # 0.106 strength = 1
     rows, cols = image.shape[:2]
     max_dim = max(rows, cols)
     
-    #Normalize the positions
-    # Y Needs to be flipped
     center_y = int(center[1] * rows)
     center_x = int(center[0] * cols)
-
-    # Inverts the Y axis (Numpy is 0 index at top of image)
     center_y = abs(rows - center_y)
 
-    print()
-    print(rows, cols)
-    print("y =", center_y, "/", rows)
-    print("x =", center_x, "/", cols)
-    print()
-    
-    pixel_radius = int(max_dim * radius)
+    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
     
-    # Calculate distance from center
     dist_from_center = np.sqrt(x**2 + y**2)
     
-    # Calculate function values
     z = func(x, y)
+    print(f"Function output - min: {np.min(z)}, max: {np.max(z)}")
     
-    # Calculate gradients
     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)}")
 
-    # Creating a sigmoid function to apply to masks
-    def sigmoid(x, center, steepness):
-        return 1 / (1 + np.exp(-steepness * (x - center)))
-    
-    print(radius)
-    print(strength)
-    print(edge_smoothness)
-    print(center_smoothness)
+    # Avoid division by zero
+    edge_smoothness = np.maximum(edge_smoothness, 1e-6)
+    center_smoothness = np.maximum(center_smoothness, 1e-6)
 
-    # Masking
     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
     
-    # 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
+    scale_factor = strength * np.log(max_dim) / 100
     gx = gx * scale_factor * mask
     gy = gy * scale_factor * mask
     
-    # Create the mapping
+    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
     
-    # 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])]
+    # Inverse transformation
+    x_inv = x - gx
+    y_inv = y - gy
     
-    transformed_image = np.dstack(channels).astype(image.dtype)
+    x_inv = x_inv * max_dim + center_x
+    y_inv = y_inv * max_dim + center_y
     
-    return transformed_image, (gx, gy)
-
+    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):
     """
@@ -142,13 +137,83 @@ def create_gradient_vector_field(gx, gy, image_shape, step=20, reverse=False):
     
     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
 #############################
 
-# pipeline = pipeline(task="image-classification", model="nick-leland/distortionml")
-
 # Version Check 
 print(f"NumPy version: {np.__version__}")
 print(f"PyTorch version: {torch.__version__}")
@@ -157,6 +222,8 @@ 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))    
@@ -177,16 +244,11 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
         return r * np.sin(theta - frequency * r)
 
     rng = np.random.default_rng()
-    if randomization_check == True:
-        radius, location, strength, edge_smoothness= definitions(rng)
-        center_x = location[0]
-        center_y = location[1]
-            
-    # Temporarily disabling and using these values.
-    # edge_smoothness = 0.25 * strength
-    # center_smoothness = 0.25 * strength
-    edge_smoothness, center_smoothness = smooth(rng, strength)
-
+    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
@@ -202,29 +264,47 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
         func = lambda x, y: spiral(x, y, frequency=spiral_frequency)
 
 
-    transformed, (gx, gy) = apply_vector_field_transform(I, func, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
-    vector_field = create_gradient_vector_field(gx, gy, I.shape[:2], reverse=reverse_gradient)
+    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)
 
-    # GRADIO CHANGE HERE
-    # predictions = pipeline(transformed) 
-
-    # Have to convert to image first
     result = Image.fromarray(transformed)
 
     categories = ['Distorted', 'Maze']
 
     def clean_output(result_values):
-        pred, idx, probs = result_values[0], result_values[1], result_values[2]
+        pred, idx, probs = result_values
         return dict(zip(categories, map(float, probs)))
 
     result_bias = learn_bias.predict(result)
     result_fresh = learn_fresh.predict(result)
-    print("Results")
     result_bias_final = clean_output(result_bias)
     result_fresh_final = clean_output(result_fresh)
 
-    # return transformed, result_bias, result_fresh, vector_field
-    return transformed, result_bias_final, result_fresh_final, vector_field
+    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,
@@ -248,10 +328,11 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Image(label="Transformed Image"),
-        # gr.Image(label="Result", num_top_classes=2)
         gr.Label(),
         gr.Label(),
-        gr.Image(label="Gradient Vector Field")
+        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)",

temp_app.py

@@ -1,286 +0,0 @@
-import numpy as np
-import gradio as gr
-from PIL import Image
-from scipy import ndimage
-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):
-    # 0.106 strength = .50
-    # 0.106 strength = 1
-    rows, cols = image.shape[:2]
-    max_dim = max(rows, cols)
-    
-    #Normalize the positions
-    # Y Needs to be flipped
-    center_y = int(center[1] * rows)
-    center_x = int(center[0] * cols)
-
-    # Inverts the Y axis (Numpy is 0 index at top of image)
-    center_y = abs(rows - center_y)
-
-    print()
-    print(rows, cols)
-    print("y =", center_y, "/", rows)
-    print("x =", center_x, "/", cols)
-    print()
-    
-    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)
-
-    # Creating a sigmoid function to apply to masks
-    def sigmoid(x, center, steepness):
-        return 1 / (1 + np.exp(-steepness * (x - center)))
-    
-    print(radius)
-    print(strength)
-    print(edge_smoothness)
-    print(center_smoothness)
-
-    # Masking
-    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
-    
-    # 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
-
-
-#############################
-#    MAIN FUNCTION HERE
-#############################
-
-# pipeline = pipeline(task="image-classification", model="nick-leland/distortionml")
-
-# 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 bulge_inverse(x, y, f=bulge, a=1, b=1, c=1, d=0, e=0):
-        t = np.arctan2(y, x)
-        term = ((f - e) / (-a))**2 - d
-        if term < 0:
-            return None, None
-
-        x = (1/np.sqrt(b)) * np.sqrt(term) * np.cos(t)
-        y = (1/np.sqrt(c)) * np.sqrt(term) * np.sin(t)
-
-        return x, y
-
-    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 == True:
-        radius, location, strength, edge_smoothness= definitions(rng)
-        center_x = location[0]
-        center_y = location[1]
-            
-    # Temporarily disabling and using these values.
-    # edge_smoothness = 0.25 * strength
-    # center_smoothness = 0.25 * strength
-    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
-        func2 = bulge_inverse
-        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)
-
-
-    # Original Image Transformation
-    transformed, (gx, gy) = apply_vector_field_transform(I, func, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
-    vector_field = create_gradient_vector_field(gx, gy, I.shape[:2], reverse=reverse_gradient)
-
-    reverted, (gx_inverse, gy_inverse) = apply_vector_field_transform(I, func2, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
-    vector_field_reverted = create_gradient_vector_field(gx_inverse, gy_inverse, I.shape[:2], reverse=reverse_gradient)
-    
-
-    # GRADIO CHANGE HERE
-    # predictions = pipeline(transformed) 
-
-    # Have to convert to image first
-    result = Image.fromarray(transformed)
-
-    categories = ['Distorted', 'Maze']
-
-    def clean_output(result_values):
-        pred, idx, probs = result_values[0], result_values[1], result_values[2]
-        return dict(zip(categories, map(float, probs)))
-
-    result_bias = learn_bias.predict(result)
-    result_fresh = learn_fresh.predict(result)
-    print("Results")
-    result_bias_final = clean_output(result_bias)
-    result_fresh_final = clean_output(result_fresh)
-
-    print("saving?")
-    result_localization = model.predict(transformed, save=True)
-    print(result_localization)
-
-    return transformed, result_bias_final, result_fresh_final, vector_field, vector_field_reverted
-
-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.Image(label="Result", num_top_classes=2)
-        gr.Label(),
-        gr.Label(),
-        gr.Image(label="Gradient Vector Field"),
-        gr.Image(label="Gradient Vector Field Reverted")
-    ],
-    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)