Update app.py

app.py

@@ -141,45 +141,6 @@ def create_gradient_vector_field(gx, gy, image_shape, step=20, reverse=False):
     
     return vector_field
 
-# 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.
@@ -297,7 +258,9 @@ def predict_image(img, model, conf_threshold, iou_threshold):
     return im
 
 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))    
+    with Image.open(image) as img:
+        img = img.convert('RGB')
+        I = np.array(img)  
 
     # Downsample large images
     max_size = 640 # Increased from 512 to allow for more detail, decreased from 1024 to match YOLO model training.
@@ -424,7 +387,7 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
         vector_field = np.zeros_like(I)
         inverted_vector_field = np.zeros_like(I)
 
-    result = Image.fromarray(transformed)
+    result = Image.fromarray(transformed.astype('uint8'), 'RGB')
 
     # categories = ['Distorted', 'Maze']