Updated the app for demo purposes

app.py

@@ -12,88 +12,88 @@ from fastai.vision.all import *
 from ultralytics import ASSETS, 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)
-    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 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)
+#     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):
     """
@@ -334,6 +334,8 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
     # strength = strength * 2  # This allows for stronger effects
 
     try:
+        strength = 0.8
+
         # Generate gradients
         gx, gy = generate_function_gradient(func, I.shape, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
         
@@ -358,16 +360,20 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
         
         inverse_transformed = cv2.remap(I, x_inv, y_inv, cv2.INTER_LINEAR)
 
+        # Apply Inverse to detected location
+        YOLO_image = predict_image(transformed, model_bulge, 0.5, 0.5)
+
+
         applied_transformed = cv2.remap(transformed, 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"Transformed image shape: {transformed.shape}")
+        # print(f"Inverse transformed image shape: {inverse_transformed.shape}")
         
         vector_field = create_gradient_vector_field(gx, gy, I.shape[:2], reverse=reverse_gradient)
         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}")
+        # 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]):
@@ -409,7 +415,6 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
     # print(result_localization2, "modelv8x")
 
 
-    YOLO_image = predict_image(transformed, model_bulge, 0.5, 0.5)
     # YOLO_image1 = predict_image(transformed, modelv8n, 0.5, 0.5)
     # YOLO_image2 = predict_image(transformed, modelv8x, 0.5, 0.5)
 
@@ -422,7 +427,7 @@ demo = gr.Interface(
     fn=transform_image,
     inputs=[
         gr.Image(type="filepath"),
-        gr.Dropdown(["Pinch", "Spiral", "Shift Up", "Bulge", "Volcano"], value="Volcano", label="Function"), 
+        gr.Dropdown(["Pinch", "Spiral", "Shift Up", "Bulge", "Volcano"], value="Bulge", 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"),
@@ -452,7 +457,7 @@ demo = gr.Interface(
     ],
     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!"
+    description=""
 )
 
 demo.launch(share=True)