Spaces:
Build error
Build error
| 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 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 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 | |
| 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 | |
| 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 | |
| ############################# | |
| # 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_bulge = YOLO("best.onnx") | |
| # modelv8x = YOLO("yolov8x.pt") | |
| # modelv8n = YOLO("yolov8n.pt") | |
| def predict_image(img, model, conf_threshold, iou_threshold): | |
| """Predicts objects in an image using a YOLOv8 model with adjustable confidence and IOU thresholds.""" | |
| results = model.predict( | |
| source=img, | |
| conf=conf_threshold, | |
| iou=iou_threshold, | |
| show_labels=True, | |
| show_conf=True, | |
| imgsz=640, | |
| ) | |
| for r in results: | |
| im_array = r.plot() | |
| im = Image.fromarray(im_array[..., ::-1]) | |
| return im | |
| def transform_image(image, func_choice, randomization_check, radius, center_x, center_y, strength, reverse_gradient=True, spiral_frequency=1): | |
| 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. | |
| 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): | |
| 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 | |
| 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: | |
| strength = 0.8 | |
| # 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) | |
| # 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}") | |
| 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}") | |
| # 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) | |
| applied_transformed = cv2.resize(applied_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() | |
| 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.astype('uint8'), 'RGB') | |
| # categories = ['Distorted', 'Maze'] | |
| # def clean_output(result_values): | |
| # pred, idx, probs = result_values | |
| # return dict(zip(categories, map(float, probs))) | |
| # Outdated, changing to a classification basis | |
| # 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_bulge.predict(transformed, save=True) | |
| print(result_localization, "bulge") | |
| # result_localization1 = modelv8n.predict(transformed, save=True) | |
| # print(result_localization1, "modelv8n") | |
| # result_localization2 = modelv8x.predict(transformed, save=True) | |
| # print(result_localization2, "modelv8x") | |
| # YOLO_image1 = predict_image(transformed, modelv8n, 0.5, 0.5) | |
| # YOLO_image2 = predict_image(transformed, modelv8x, 0.5, 0.5) | |
| # return transformed, YOLO_image, YOLO_image1, YOLO_image2, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field | |
| # return transformed, YOLO_image, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field | |
| return transformed, YOLO_image, vector_field, inverse_transformed, inverted_vector_field, applied_transformed | |
| demo = gr.Interface( | |
| fn=transform_image, | |
| inputs=[ | |
| gr.Image(type="filepath"), | |
| 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"), | |
| 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="bulge_model Model Classification"), | |
| # gr.Image(label="yolov8n Model Classification"), | |
| # gr.Image(label="yolov8x Model Classification"), | |
| # gr.Label(), | |
| # gr.Label(), | |
| gr.Image(label="Gradient Vector Field"), | |
| gr.Image(label="Inverse Gradient"), | |
| gr.Image(label="Inverted Vector Field"), | |
| gr.Image(label="Fixed Image") | |
| ], | |
| 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="" | |
| ) | |
| demo.launch(share=True) | |