Updated the app, reformated the way that the app functions

app.py

@@ -10,10 +10,14 @@ import fastai
 from fastcore.all import *
 from fastai.vision.all import *
 from ultralytics import YOLO
+import cv2
 
 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)
+    print()
+    print(f"Max_dim is {max_dim}")
+    print()
     
     center_y = int(center[1] * rows)
     center_x = int(center[0] * cols)
@@ -137,9 +141,6 @@ 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.
@@ -210,6 +211,57 @@ def apply_gradient_transform(image, gx, gy):
     
     return transformed_image
 
+def generate_function_gradient(func, image_shape, 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)
+    
+    y, x = np.mgrid[0:rows, 0:cols].astype(np.float32)
+    y = (y - center[1] * rows) / max_dim
+    x = (x - center[0] * cols) / max_dim
+    
+    dist_from_center = np.sqrt(x**2 + y**2)
+    
+    z = func(x, y)
+    
+    gy, gx = np.gradient(z)
+
+    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 *= mask
+    gy *= mask
+    
+    magnitude = np.sqrt(gx**2 + gy**2)
+    max_magnitude = np.max(magnitude)
+    if max_magnitude > 0:
+        gx /= max_magnitude
+        gy /= max_magnitude
+    
+    # Increase the base scale factor
+    base_scale = radius * max_dim * 0.2  # Increased from 0.1 to 0.2
+    
+    # Apply a non-linear scaling to the strength
+    adjusted_strength = np.power(strength, 1.5)  # This will make the effect more pronounced at higher strengths
+    
+    # Increase the maximum strength multiplier
+    scale_factor = base_scale * np.clip(adjusted_strength, 0, 3)  # Increased max from 2 to 3
+    
+    # Apply an additional scaling factor based on image size
+    size_factor = np.log(max_dim) / np.log(1000)  # This will be 1 for 1000x1000 images, larger for bigger images
+    scale_factor *= size_factor
+    
+    gx *= scale_factor
+    gy *= scale_factor
+    
+    print(f"Final scale factor: {scale_factor}")
+    print(f"Final gradient ranges: gx [{np.min(gx)}, {np.max(gx)}], gy [{np.min(gy)}, {np.max(gy)}]")
+    
+    return gx, gy
+
 #############################
 #    MAIN FUNCTION HERE
 #############################
@@ -228,7 +280,23 @@ 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))    
 
+    # Downsample large images
+    max_size = 1024  # Increased from 512 to allow for more detail
+    if max(I.shape[:2]) > max_size:
+        scale = max_size / max(I.shape[:2])
+        new_size = (int(I.shape[1] * scale), int(I.shape[0] * scale))
+        I = cv2.resize(I, new_size, interpolation=cv2.INTER_AREA)
+        print(f"Downsampled image to {I.shape}")
+
+    ##################################
+    #    Transformation Functions    #
+    ##################################
+
     def pinch(x, y):
+        r = np.sqrt(x**2 + y**2)
+        return r
+
+    def zoom(x, y):
         return x**2 + y**2
 
     def shift(x, y):
@@ -247,40 +315,82 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
     if randomization_check:
         radius, location, strength, edge_smoothness = definitions(rng)
         center_x, center_y = location
+        center_smoothness = edge_smoothness
     else:
         edge_smoothness, center_smoothness = smooth(rng, strength)
 
     if func_choice == "Pinch":
         func = pinch
+        edge_smoothness = 0
+        center_smoothness = 0
+
     elif func_choice == "Spiral":
         func = shift 
+        edge_smoothness = 0
+        center_smoothness = 0
+
     elif func_choice == "Bulge":
         func = bulge
         edge_smoothness = 0
         center_smoothness = 0
     elif func_choice == "Volcano":
         func = bulge
+        edge_smoothness = 0
+        center_smoothness = 0
+
     elif func_choice == "Shift Up":
         func = lambda x, y: spiral(x, y, frequency=spiral_frequency)
+        edge_smoothness = 0
+        center_smoothness = 0
 
 
     print(f"Function choice: {func_choice}")
     print(f"Input image shape: {I.shape}")
+    print(f"Radius: {radius}, Center: ({center_x}, {center_y}), Strength: {strength}")
+
+    # strength = strength * 2  # This allows for stronger effects
 
     try:
-        transformed, inverse_transformed, (gx, gy) = apply_vector_field_transform(
-            I, func, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness
-        )
+        # Generate gradients
+        gx, gy = generate_function_gradient(func, I.shape, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
+        
+        # Vectorized transformation
+        rows, cols = I.shape[:2]
+        y, x = np.mgrid[0:rows, 0:cols].astype(np.float32)
+        
+        x_new = x + gx
+        y_new = y + gy
+        
+        x_new = np.clip(x_new, 0, cols - 1)
+        y_new = np.clip(y_new, 0, rows - 1)
+        
+        transformed = cv2.remap(I, x_new, y_new, cv2.INTER_LINEAR)
+        
+        inv_gx, inv_gy = -gx, -gy
+        x_inv = x + inv_gx
+        y_inv = y + inv_gy
+        x_inv = np.clip(x_inv, 0, cols - 1)
+        y_inv = np.clip(y_inv, 0, rows - 1)
+        
+        inverse_transformed = cv2.remap(I, x_inv, y_inv, cv2.INTER_LINEAR)
+
         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)
+        inverted_vector_field = create_gradient_vector_field(inv_gx, inv_gy, I.shape[:2], reverse=False)
         
         print(f"Vector field shape: {vector_field.shape}")
         print(f"Inverted vector field shape: {inverted_vector_field.shape}")
+
+        # If we downsampled earlier, upsample the results back to original size
+        if max(I.shape[:2]) != max(np.asarray(Image.open(image)).shape[:2]):
+            original_size = np.asarray(Image.open(image)).shape[:2][::-1]
+            transformed = cv2.resize(transformed, original_size, interpolation=cv2.INTER_LINEAR)
+            inverse_transformed = cv2.resize(inverse_transformed, original_size, interpolation=cv2.INTER_LINEAR)
+            vector_field = cv2.resize(vector_field, original_size, interpolation=cv2.INTER_LINEAR)
+            inverted_vector_field = cv2.resize(inverted_vector_field, original_size, interpolation=cv2.INTER_LINEAR)
+
     except Exception as e:
         print(f"Error in transformation: {str(e)}")
         traceback.print_exc()
@@ -306,6 +416,7 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
 
     return transformed, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field
 
+
 demo = gr.Interface(
     fn=transform_image,
     inputs=[