nick-leland commited on
Commit
8093ae7
·
1 Parent(s): e5010cc

Worked on the other aspects

Browse files
Files changed (2) hide show
  1. temp_app.py +286 -0
  2. test.py +116 -0
temp_app.py ADDED
@@ -0,0 +1,286 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import gradio as gr
3
+ from PIL import Image
4
+ from scipy import ndimage
5
+ import matplotlib.pyplot as plt
6
+ from bulk_bulge_generation import definitions, smooth
7
+ # from transformers import pipeline
8
+ import fastai
9
+ from fastcore.all import *
10
+ from fastai.vision.all import *
11
+ from ultralytics import YOLO
12
+
13
+ def apply_vector_field_transform(image, func, radius, center=(0.5, 0.5), strength=1, edge_smoothness=0.1, center_smoothness=0.20):
14
+ # 0.106 strength = .50
15
+ # 0.106 strength = 1
16
+ rows, cols = image.shape[:2]
17
+ max_dim = max(rows, cols)
18
+
19
+ #Normalize the positions
20
+ # Y Needs to be flipped
21
+ center_y = int(center[1] * rows)
22
+ center_x = int(center[0] * cols)
23
+
24
+ # Inverts the Y axis (Numpy is 0 index at top of image)
25
+ center_y = abs(rows - center_y)
26
+
27
+ print()
28
+ print(rows, cols)
29
+ print("y =", center_y, "/", rows)
30
+ print("x =", center_x, "/", cols)
31
+ print()
32
+
33
+ pixel_radius = int(max_dim * radius)
34
+
35
+ y, x = np.ogrid[:rows, :cols]
36
+ y = (y - center_y) / max_dim
37
+ x = (x - center_x) / max_dim
38
+
39
+ # Calculate distance from center
40
+ dist_from_center = np.sqrt(x**2 + y**2)
41
+
42
+ # Calculate function values
43
+ z = func(x, y)
44
+
45
+ # Calculate gradients
46
+ gy, gx = np.gradient(z)
47
+
48
+ # Creating a sigmoid function to apply to masks
49
+ def sigmoid(x, center, steepness):
50
+ return 1 / (1 + np.exp(-steepness * (x - center)))
51
+
52
+ print(radius)
53
+ print(strength)
54
+ print(edge_smoothness)
55
+ print(center_smoothness)
56
+
57
+ # Masking
58
+ edge_mask = np.clip((radius - dist_from_center) / (radius * edge_smoothness), 0, 1)
59
+
60
+ center_mask = np.clip((dist_from_center - radius * center_smoothness) / (radius * center_smoothness), 0, 1)
61
+
62
+ mask = edge_mask * center_mask
63
+
64
+ # Apply mask to gradients
65
+ gx = gx * mask
66
+ gy = gy * mask
67
+
68
+ # Normalize gradient vectors
69
+ magnitude = np.sqrt(gx**2 + gy**2)
70
+ magnitude[magnitude == 0] = 1 # Avoid division by zero
71
+ gx = gx / magnitude
72
+ gy = gy / magnitude
73
+
74
+ # Scale the effect (Play with the number 5)
75
+ scale_factor = strength * np.log(max_dim) / 100 # Adjust strength based on image size
76
+ gx = gx * scale_factor * mask
77
+ gy = gy * scale_factor * mask
78
+
79
+ # Create the mapping
80
+ x_new = x + gx
81
+ y_new = y + gy
82
+
83
+ # Convert back to pixel coordinates
84
+ x_new = x_new * max_dim + center_x
85
+ y_new = y_new * max_dim + center_y
86
+
87
+ # Ensure the new coordinates are within the image boundaries
88
+ x_new = np.clip(x_new, 0, cols - 1)
89
+ y_new = np.clip(y_new, 0, rows - 1)
90
+
91
+ # Apply the transformation to each channel
92
+ channels = [ndimage.map_coordinates(image[..., i], [y_new, x_new], order=1, mode='reflect')
93
+ for i in range(image.shape[2])]
94
+
95
+ transformed_image = np.dstack(channels).astype(image.dtype)
96
+
97
+ return transformed_image, (gx, gy)
98
+
99
+
100
+ def create_gradient_vector_field(gx, gy, image_shape, step=20, reverse=False):
101
+ """
102
+ Create a gradient vector field visualization with option to reverse direction.
103
+
104
+ :param gx: X-component of the gradient
105
+ :param gy: Y-component of the gradient
106
+ :param image_shape: Shape of the original image (height, width)
107
+ :param step: Spacing between arrows
108
+ :param reverse: If True, reverse the direction of the arrows
109
+ :return: Gradient vector field as a numpy array (RGB image)
110
+ """
111
+ rows, cols = image_shape
112
+ y, x = np.mgrid[step/2:rows:step, step/2:cols:step].reshape(2, -1).astype(int)
113
+
114
+ # Calculate the scale based on image size
115
+ max_dim = max(rows, cols)
116
+ scale = max_dim / 1000 # Adjusted for longer arrows
117
+
118
+ # Reverse direction if specified
119
+ direction = -1 if reverse else 1
120
+
121
+ fig, ax = plt.subplots(figsize=(cols/50, rows/50), dpi=100)
122
+ ax.quiver(x, y, direction * gx[y, x], direction * -gy[y, x],
123
+ scale=scale,
124
+ scale_units='width',
125
+ width=0.002 * max_dim / 500,
126
+ headwidth=8,
127
+ headlength=12,
128
+ headaxislength=0,
129
+ color='black',
130
+ minshaft=2,
131
+ minlength=0,
132
+ pivot='tail')
133
+ ax.set_xlim(0, cols)
134
+ ax.set_ylim(rows, 0)
135
+ ax.set_aspect('equal')
136
+ ax.axis('off')
137
+
138
+ fig.tight_layout(pad=0)
139
+ fig.canvas.draw()
140
+ vector_field = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
141
+ vector_field = vector_field.reshape(fig.canvas.get_width_height()[::-1] + (3,))
142
+ plt.close(fig)
143
+
144
+ return vector_field
145
+
146
+
147
+ #############################
148
+ # MAIN FUNCTION HERE
149
+ #############################
150
+
151
+ # pipeline = pipeline(task="image-classification", model="nick-leland/distortionml")
152
+
153
+ # Version Check
154
+ print(f"NumPy version: {np.__version__}")
155
+ print(f"PyTorch version: {torch.__version__}")
156
+ print(f"FastAI version: {fastai.__version__}")
157
+
158
+ learn_bias = load_learner('model_bias.pkl')
159
+ learn_fresh = load_learner('model_fresh.pkl')
160
+
161
+ # Loads the YOLO Model
162
+ model = YOLO("bulge_yolo_model.pt")
163
+
164
+
165
+ def transform_image(image, func_choice, randomization_check, radius, center_x, center_y, strength, reverse_gradient=True, spiral_frequency=1):
166
+ I = np.asarray(Image.open(image))
167
+
168
+ def pinch(x, y):
169
+ return x**2 + y**2
170
+
171
+ def shift(x, y):
172
+ return np.arctan2(y, x)
173
+
174
+ def bulge(x, y):
175
+ r = -np.sqrt(x**2 + y**2)
176
+ return r
177
+
178
+ def bulge_inverse(x, y, f=bulge, a=1, b=1, c=1, d=0, e=0):
179
+ t = np.arctan2(y, x)
180
+ term = ((f - e) / (-a))**2 - d
181
+ if term < 0:
182
+ return None, None
183
+
184
+ x = (1/np.sqrt(b)) * np.sqrt(term) * np.cos(t)
185
+ y = (1/np.sqrt(c)) * np.sqrt(term) * np.sin(t)
186
+
187
+ return x, y
188
+
189
+ def spiral(x, y, frequency=1):
190
+ r = np.sqrt(x**2 + y**2)
191
+ theta = np.arctan2(y, x)
192
+ return r * np.sin(theta - frequency * r)
193
+
194
+ rng = np.random.default_rng()
195
+ if randomization_check == True:
196
+ radius, location, strength, edge_smoothness= definitions(rng)
197
+ center_x = location[0]
198
+ center_y = location[1]
199
+
200
+ # Temporarily disabling and using these values.
201
+ # edge_smoothness = 0.25 * strength
202
+ # center_smoothness = 0.25 * strength
203
+ edge_smoothness, center_smoothness = smooth(rng, strength)
204
+
205
+
206
+ if func_choice == "Pinch":
207
+ func = pinch
208
+ elif func_choice == "Spiral":
209
+ func = shift
210
+ elif func_choice == "Bulge":
211
+ func = bulge
212
+ func2 = bulge_inverse
213
+ edge_smoothness = 0
214
+ center_smoothness = 0
215
+ elif func_choice == "Volcano":
216
+ func = bulge
217
+ elif func_choice == "Shift Up":
218
+ func = lambda x, y: spiral(x, y, frequency=spiral_frequency)
219
+
220
+
221
+ # Original Image Transformation
222
+ transformed, (gx, gy) = apply_vector_field_transform(I, func, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
223
+ vector_field = create_gradient_vector_field(gx, gy, I.shape[:2], reverse=reverse_gradient)
224
+
225
+ reverted, (gx_inverse, gy_inverse) = apply_vector_field_transform(I, func2, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
226
+ vector_field_reverted = create_gradient_vector_field(gx_inverse, gy_inverse, I.shape[:2], reverse=reverse_gradient)
227
+
228
+
229
+ # GRADIO CHANGE HERE
230
+ # predictions = pipeline(transformed)
231
+
232
+ # Have to convert to image first
233
+ result = Image.fromarray(transformed)
234
+
235
+ categories = ['Distorted', 'Maze']
236
+
237
+ def clean_output(result_values):
238
+ pred, idx, probs = result_values[0], result_values[1], result_values[2]
239
+ return dict(zip(categories, map(float, probs)))
240
+
241
+ result_bias = learn_bias.predict(result)
242
+ result_fresh = learn_fresh.predict(result)
243
+ print("Results")
244
+ result_bias_final = clean_output(result_bias)
245
+ result_fresh_final = clean_output(result_fresh)
246
+
247
+ print("saving?")
248
+ result_localization = model.predict(transformed, save=True)
249
+ print(result_localization)
250
+
251
+ return transformed, result_bias_final, result_fresh_final, vector_field, vector_field_reverted
252
+
253
+ demo = gr.Interface(
254
+ fn=transform_image,
255
+ inputs=[
256
+ gr.Image(type="filepath"),
257
+ gr.Dropdown(["Pinch", "Spiral", "Shift Up", "Bulge", "Volcano"], value="Volcano", label="Function"),
258
+ gr.Checkbox(label="Randomize inputs?"),
259
+ gr.Slider(0, 0.5, value=0.25, label="Radius (as fraction of image size)"),
260
+ gr.Slider(0, 1, value=0.5, label="Center X"),
261
+ gr.Slider(0, 1, value=0.5, label="Center Y"),
262
+ gr.Slider(0, 1, value=0.5, label="Strength"),
263
+ # gr.Slider(0, 1, value=0.5, label="Edge Smoothness"),
264
+ # gr.Slider(0, 0.5, value=0.1, label="Center Smoothness")
265
+ # gr.Checkbox(label="Reverse Gradient Direction"),
266
+ ],
267
+ examples=[
268
+ [np.asarray(Image.open("examples/1500_maze.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
269
+ [np.asarray(Image.open("examples/2048_maze.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
270
+ [np.asarray(Image.open("examples/2300_fresh.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
271
+ [np.asarray(Image.open("examples/50_fresh.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5]
272
+ ],
273
+ outputs=[
274
+ gr.Image(label="Transformed Image"),
275
+ # gr.Image(label="Result", num_top_classes=2)
276
+ gr.Label(),
277
+ gr.Label(),
278
+ gr.Image(label="Gradient Vector Field"),
279
+ gr.Image(label="Gradient Vector Field Reverted")
280
+ ],
281
+ title="Image Transformation Demo!",
282
+ article="If you like this demo, please star the github repository for the project! Located [here!](https://github.com/nick-leland/DistortionML)",
283
+ 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!"
284
+ )
285
+
286
+ demo.launch(share=True)
test.py ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import matplotlib.pyplot as plt
3
+
4
+ def function(x, y, a=1, b=1, c=1, d=0, e=0):
5
+ return -a * np.sqrt((b * x)**2 + (c * y)**2 + d) + e
6
+
7
+ def inverse_function(f, t, a, b, c, d, e):
8
+ term = ((f - e) / (-a))**2 - d
9
+ if term < 0:
10
+ return None, None
11
+
12
+ x = (1/np.sqrt(b)) * np.sqrt(term) * np.cos(t)
13
+ y = (1/np.sqrt(c)) * np.sqrt(term) * np.sin(t)
14
+
15
+ return x, y
16
+
17
+ def create_gradient_vector_field(gx, gy, image_shape, step=20, reverse=False):
18
+ rows, cols = image_shape
19
+ y, x = np.mgrid[step/2:rows:step, step/2:cols:step].reshape(2, -1).astype(int)
20
+
21
+ max_dim = max(rows, cols)
22
+ scale = max_dim / 1000
23
+
24
+ direction = -1 if reverse else 1
25
+
26
+ fig, ax = plt.subplots(figsize=(cols/50, rows/50), dpi=100)
27
+ ax.quiver(x, y, direction * gx[y, x], direction * -gy[y, x],
28
+ scale=scale,
29
+ scale_units='width',
30
+ width=0.002 * max_dim / 500,
31
+ headwidth=8,
32
+ headlength=12,
33
+ headaxislength=0,
34
+ color='black',
35
+ minshaft=2,
36
+ minlength=0,
37
+ pivot='tail')
38
+ ax.set_xlim(0, cols)
39
+ ax.set_ylim(rows, 0)
40
+ ax.set_aspect('equal')
41
+ ax.axis('off')
42
+
43
+ fig.tight_layout(pad=0)
44
+ fig.canvas.draw()
45
+ vector_field = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
46
+ vector_field = vector_field.reshape(fig.canvas.get_width_height()[::-1] + (3,))
47
+ plt.close(fig)
48
+
49
+ return vector_field
50
+
51
+
52
+ def apply_inverse_vector_field_transform(image, func, radius, center=(0.5, 0.5), strength=1, edge_smoothness=0.1, center_smoothness=0.20):
53
+
54
+ rows, cols = image.shape[:2]
55
+ max_dim = max(rows, cols)
56
+
57
+ center_y = int(center[1] * rows)
58
+ center_x = int(center[0] * cols)
59
+ center_y = abs(rows - center_y)
60
+
61
+ pixel_radius = int(max_dim * radius)
62
+
63
+ y, x = np.ogrid[:rows, :cols]
64
+ y = (y - center_y) / max_dim
65
+ x = (x - center_x) / max_dim
66
+
67
+ dist_from_center = np.sqrt(x**2 + y**2)
68
+
69
+ z = func(x, y)
70
+ gy, gx = np.gradient(z)
71
+
72
+ def sigmoid(x, center, steepness):
73
+ return 1 / (1+ np.exp(-steepness * (x - center)))
74
+
75
+ mask = edge_mask * center_mask
76
+
77
+ gx = gx * mask
78
+ gy = gy * mask
79
+
80
+ magnitude = np.sqrt(gx**2 + gy**2)
81
+ magnitude[magnitude == 0] = 1
82
+ gx = gx / magnitude
83
+ gy = gy / magnitude
84
+
85
+ scale_factor = strength * np.log(max_dim) / 100
86
+ gx = gx * scale_factor * mask
87
+ gy = gy * scale_factor * mask
88
+
89
+ x_new = x + gx
90
+ y_new = y + gy
91
+
92
+ x_new = y_new * max_dim + center_x
93
+ y_new = y_new * max_dim + center_y
94
+
95
+ x_new = np.clip(x_new, 0, cols - 1)
96
+ y_new = np.clip(y_new, 0, rows - 1)
97
+
98
+
99
+
100
+
101
+
102
+ if __name__ == '__main__':
103
+ x = 3
104
+ y = 2
105
+
106
+ t = np.arctan2(y, x)
107
+
108
+ a, b, c, d, e = 1, 1, 1, 0, 0
109
+
110
+ print(x, y)
111
+ function = function(3, 2)
112
+ print(function)
113
+ inverse_function = inverse_function(function, t, a, b, c, d, e)
114
+ print(inverse_function)
115
+
116
+