xu3kev commited on
Commit
bcbc8c1
1 Parent(s): 65dd20f

Update myturtle.py

Browse files
Files changed (1) hide show
  1. myturtle.py +272 -10
myturtle.py CHANGED
@@ -1,34 +1,185 @@
1
  import numpy as np
2
  import cv2
3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
  HALF_INF = 63
5
  INF = 126
6
  EPS_DIST = 1/20
7
  EPS_ANGLE = 2.86
8
- SCALE = 15.5
 
 
 
 
9
 
10
  class Turtle:
11
- def __init__(self, canvas_size=(2000, 2000)):
12
- self.x = canvas_size[0] // 2
13
- self.y = canvas_size[1] // 2
14
  self.heading = 0
15
  self.canvas = np.ones((canvas_size[1], canvas_size[0], 3), dtype=np.uint8) * 255
16
  self.is_down = True
 
 
 
17
 
18
  def forward(self, dist):
 
 
 
 
19
  dist = dist * SCALE
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
  x0, y0 = self.x, self.y
21
- x1 = int(x0 + dist * np.cos(self.heading))
22
- y1 = int(y0 - dist * np.sin(self.heading))
23
  if self.is_down:
24
- cv2.line(self.canvas, (x0, y0), (x1, y1), (0, 0, 0), 3)
25
  self.x, self.y = x1, y1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26
 
27
  def left(self, angle):
28
- self.heading += angle * np.pi / 180
 
 
 
 
29
 
30
  def right(self, angle):
31
- self.heading -= angle * np.pi / 180
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
32
 
33
  def penup(self):
34
  self.is_down = False
@@ -40,6 +191,17 @@ class Turtle:
40
  if path:
41
  cv2.imwrite(path, self.canvas)
42
  return self.canvas
 
 
 
 
 
 
 
 
 
 
 
43
 
44
  class _TurtleState:
45
  def __init__(self, turtle):
@@ -113,6 +275,7 @@ if __name__ == "__main__":
113
  right(90)
114
  forward(50)
115
  save("test2.png")
 
116
 
117
  def plot2():
118
  for j in range(2):
@@ -127,5 +290,104 @@ if __name__ == "__main__":
127
  left(180.0)
128
  FINAL_IMAGE = turtle.save("")
129
 
130
- example_plot()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
131
  # plot2()
 
1
  import numpy as np
2
  import cv2
3
 
4
+
5
+ def crop_and_scaled_imgs(imgs):
6
+ PAD = 10
7
+ # use the last image to find the bounding box of the non-white area and the transformation parameters
8
+ # and then apply the transformation to all images
9
+
10
+
11
+ img = imgs[-1]
12
+ gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
13
+ # Threshold the image to create a binary mask
14
+ _, binary_mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
15
+
16
+ # Find the coordinates of non-zero pixels
17
+ coords = cv2.findNonZero(binary_mask)
18
+
19
+ # Get the bounding box of the non-zero pixels
20
+ x, y, w, h = cv2.boundingRect(coords)
21
+ x -= PAD
22
+ y -= PAD
23
+ w += 2 * PAD
24
+ h += 2 * PAD
25
+
26
+ # Calculate the position to center the ROI in the 256x256 image
27
+ start_x = max(0, (256 - w) // 2)
28
+ start_y = max(0, (256 - h) // 2)
29
+
30
+ # Create a new 256x256 rgb images
31
+ new_imgs = [np.ones((256, 256, 3), dtype=np.uint8) * 255 for _ in range(len(imgs))]
32
+ for i in range(len(imgs)):
33
+ # Extract the ROI (region of interest) of the non-white area
34
+ roi = imgs[i][y:y+h, x:x+w]
35
+ # If the ROI is larger than 256x256, resize it
36
+
37
+ if w > 256 or h > 256:
38
+ scale = min(256 / w, 256 / h)
39
+ new_w = int(w * scale)
40
+ new_h = int(h * scale)
41
+ roi = cv2.resize(roi, (new_w, new_h), interpolation=cv2.INTER_AREA)
42
+ w, h = new_w, new_h
43
+
44
+ # new_imgs[i] = np.ones((256, 256), dtype=np.uint8) * 255
45
+ # centered_img = np.ones((256, 256), dtype=np.uint8) * 255
46
+
47
+ # Place the ROI in the centered position
48
+ new_imgs[i][start_y:start_y+h, start_x:start_x+w] = roi
49
+
50
+ return new_imgs
51
+
52
+
53
  HALF_INF = 63
54
  INF = 126
55
  EPS_DIST = 1/20
56
  EPS_ANGLE = 2.86
57
+ SCALE = 15
58
+
59
+ MOVE_SPEED = 25
60
+ ROTATE_SPEED = 30
61
+ FPS = 24
62
 
63
  class Turtle:
64
+ def __init__(self, canvas_size=(800, 800)):
65
+ self.x = canvas_size[0] // 2
66
+ self.y = canvas_size[1] // 2
67
  self.heading = 0
68
  self.canvas = np.ones((canvas_size[1], canvas_size[0], 3), dtype=np.uint8) * 255
69
  self.is_down = True
70
+ self.time_since_last_frame = 0
71
+ self.frames = [self.canvas.copy()]
72
+
73
 
74
  def forward(self, dist):
75
+ # print('st', self.x, self.y)
76
+ # self.forward_step(dist * SCALE)
77
+ # print('ed', self.x, self.y)
78
+ # return
79
  dist = dist * SCALE
80
+ sign = 1 if dist > 0 else -1
81
+ abs_dist = abs(dist)
82
+ if self.time_since_last_frame + abs_dist / MOVE_SPEED >= 1:
83
+ dist1 = (1 - self.time_since_last_frame) * MOVE_SPEED
84
+ self.forward_step(dist1 * sign)
85
+ self.save_frame_with_turtle()
86
+ self.time_since_last_frame = 0
87
+ # for loop to step forward
88
+ num_steps = int((abs_dist - dist1) / MOVE_SPEED)
89
+ for _ in range(num_steps):
90
+ self.forward_step(MOVE_SPEED * sign)
91
+ self.save_frame_with_turtle()
92
+ last_abs_dist = abs_dist - dist1 - num_steps * MOVE_SPEED
93
+ if last_abs_dist >= MOVE_SPEED:
94
+ self.forward_step(MOVE_SPEED * sign)
95
+ self.save_frame_with_turtle()
96
+ last_abs_dist -= MOVE_SPEED
97
+ self.forward_step(last_abs_dist * sign)
98
+ self.time_since_last_frame = last_abs_dist / MOVE_SPEED
99
+ else:
100
+ self.forward_step(abs_dist * sign)
101
+ # self.time_since_last_frame += abs_dist / MOVE_SPEED
102
+ # if self.time_since_last_frame >= 1:
103
+ # self.time_since_last_frame = 0
104
+
105
+ def forward_step(self, dist):
106
+ # print('step', dist)
107
+ if dist == 0:
108
+ return
109
  x0, y0 = self.x, self.y
110
+ x1 = (x0 + dist * np.cos(self.heading))
111
+ y1 = (y0 - dist * np.sin(self.heading))
112
  if self.is_down:
113
+ cv2.line(self.canvas, (int(np.rint(x0)), int(np.rint(y0))), (int(np.rint(x1)), int(np.rint(y1))), (0, 0, 0), 3)
114
  self.x, self.y = x1, y1
115
+ self.time_since_last_frame += abs(dist) / MOVE_SPEED
116
+ # self.frames.append(self.canvas.copy())
117
+ # self.save_frame_with_turtle()
118
+
119
+ def save_frame_with_turtle(self):
120
+ # save the current frame to frames buffer
121
+ # also plot a red triangle to represent the turtle pointing to the current direction
122
+
123
+ # draw the turtle
124
+ x, y = self.x, self.y
125
+ canvas_copy = self.canvas.copy()
126
+ triangle_size = 10
127
+ x0 = int(np.rint(x + triangle_size * np.cos(self.heading)))
128
+ y0 = int(np.rint(y - triangle_size * np.sin(self.heading)))
129
+ x1 = int(np.rint(x + triangle_size * np.cos(self.heading + 2 * np.pi / 3)))
130
+ y1 = int(np.rint(y - triangle_size * np.sin(self.heading + 2 * np.pi / 3)))
131
+ x2 = int(np.rint(x + triangle_size * np.cos(self.heading - 2 * np.pi / 3)))
132
+ y2 = int(np.rint(y - triangle_size * np.sin(self.heading - 2 * np.pi / 3)))
133
+ x3 = int(np.rint(x - 0.25 * triangle_size * np.cos(self.heading)))
134
+ y3 = int(np.rint(y + 0.25 * triangle_size * np.sin(self.heading)))
135
+ # fill the triangle
136
+ cv2.fillPoly(canvas_copy, [np.array([(x0, y0), (x1, y1), (x3, y3), (x2, y2)], dtype=np.int32)], (0, 0, 255))
137
+
138
+ self.frames.append(canvas_copy)
139
+
140
+
141
 
142
  def left(self, angle):
143
+ # print('angel', angle)
144
+ # print('ast', self.heading)
145
+ # self.heading += angle * np.pi / 180
146
+ self.turn_to(angle)
147
+ # print('aed', self.heading)
148
 
149
  def right(self, angle):
150
+ # print('angel', angle)
151
+ # print('ast', self.heading)
152
+ # self.heading -= angle * np.pi / 180
153
+ self.turn_to(-angle)
154
+ # print('aed', self.heading)
155
+
156
+ def turn_to(self, angle):
157
+ abs_angle = abs(angle)
158
+ sign = 1 if angle > 0 else -1
159
+ if self.time_since_last_frame + abs(angle) / ROTATE_SPEED > 1:
160
+ angle1 = (1 - self.time_since_last_frame) * ROTATE_SPEED
161
+ self.turn_to_step(angle1 * sign)
162
+ self.save_frame_with_turtle()
163
+ self.time_since_last_frame = 0
164
+ num_steps = int((abs_angle - angle1) / ROTATE_SPEED)
165
+ for _ in range(num_steps):
166
+ self.turn_to_step(ROTATE_SPEED * sign)
167
+ self.save_frame_with_turtle()
168
+ last_abs_angle = abs_angle - angle1 - num_steps * ROTATE_SPEED
169
+ if last_abs_angle >= ROTATE_SPEED:
170
+ self.turn_to_step(ROTATE_SPEED * sign)
171
+ self.save_frame_with_turtle()
172
+ last_abs_angle -= ROTATE_SPEED
173
+ self.turn_to_step(last_abs_angle * sign)
174
+ self.time_since_last_frame = last_abs_angle / ROTATE_SPEED
175
+ else:
176
+ self.turn_to_step(abs_angle * sign)
177
+ # self.time_since_last_frame += abs_angle / ROTATE_SPEED
178
+
179
+ def turn_to_step(self, angle):
180
+ # print('turn step', angle)
181
+ self.heading += angle * np.pi / 180
182
+ self.time_since_last_frame += abs(angle) / ROTATE_SPEED
183
 
184
  def penup(self):
185
  self.is_down = False
 
191
  if path:
192
  cv2.imwrite(path, self.canvas)
193
  return self.canvas
194
+
195
+ def save_gif(self, path):
196
+ import imageio.v3 as iio
197
+ frames_rgb = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in self.frames]
198
+ print(f'number of frames: {len(frames_rgb)}')
199
+ frames_rgb.extend(FPS*2 * [frames_rgb[-1]])
200
+
201
+ frames_rgb = crop_and_scaled_imgs(frames_rgb)
202
+ # iio.imwrite(path, np.stack(frames_rgb), fps=30, plugin='pillow')
203
+ return iio.imwrite('<bytes>', np.stack(frames_rgb), fps=FPS, loop=0, plugin='pillow', format='gif')
204
+
205
 
206
  class _TurtleState:
207
  def __init__(self, turtle):
 
275
  right(90)
276
  forward(50)
277
  save("test2.png")
278
+ return turtle.frames
279
 
280
  def plot2():
281
  for j in range(2):
 
290
  left(180.0)
291
  FINAL_IMAGE = turtle.save("")
292
 
293
+ def plot3():
294
+ frames = []
295
+ frames.append(np.array(turtle.save("")))
296
+ for j in range(2):
297
+ forward(2)
298
+ frames.append(np.array(turtle.save("")))
299
+ left(0.0)
300
+ for i in range(4):
301
+ forward(2)
302
+ left(90)
303
+ frames.append(np.array(turtle.save("")))
304
+ forward(0)
305
+ left(180.0)
306
+ forward(2)
307
+ left(180.0)
308
+ frames.append(np.array(turtle.save("")))
309
+
310
+ return frames
311
+
312
+ def make_gif(frames, filename):
313
+ import imageio
314
+ frames_rgb = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in frames]
315
+ imageio.mimsave(filename, frames_rgb, fps=30)
316
+
317
+ def make_gif2(frames, filename):
318
+ import imageio.v3 as iio
319
+ frames_rgb = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in frames]
320
+ print(f'number of frames: {len(frames_rgb)}')
321
+ iio.imwrite(filename, np.stack(frames_rgb), fps=30, plugin='pillow')
322
+
323
+ def make_gif3(frames, filename):
324
+ from moviepy.editor import ImageSequenceClip
325
+ clip = ImageSequenceClip(list(frames), fps=20)
326
+ clip.write_gif(filename, fps=20)
327
+
328
+ def make_gif4(frames, filename):
329
+ from array2gif import write_gif
330
+ write_gif(frames, filename, fps=20)
331
+
332
+ def make_gif5(frames, filename):
333
+ from PIL import Image
334
+ frames_rgb = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in frames]
335
+ images = [Image.fromarray(frame) for frame in frames_rgb]
336
+ images[0].save(filename, save_all=True, append_images=images[1:], duration=100, loop=0)
337
+
338
+
339
+
340
+ def plot4():
341
+ # the following program draws a treelike pattern
342
+ import random
343
+
344
+ def draw_tree(level, length, angle):
345
+ if level == 0:
346
+ return
347
+ else:
348
+ forward(length)
349
+ left(angle)
350
+ draw_tree(level-1, length*0.7, angle*0.8)
351
+ right(angle*2)
352
+ draw_tree(level-1, length*0.7, angle*0.8)
353
+ left(angle)
354
+ forward(-length)
355
+
356
+ random.seed(0) # Comment this line to change the randomness
357
+ for _ in range(7): # Adjust the number to control the density
358
+ draw_tree(5, 5, 30)
359
+ forward(0)
360
+ left(random.randint(0, 360))
361
+ turtle.save("test3.png")
362
+ return turtle.frames
363
+
364
+ def plot5():
365
+ for i in range(7):
366
+ with fork_state():
367
+ for j in range(4):
368
+ forward(2*i)
369
+ left(90.0)
370
+ return turtle.frames
371
+
372
+
373
+ # make_gif2(plot5(), "test.gif")
374
+ frames = plot5()
375
+ # frames = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in frames]
376
+ # breakpoint()
377
+ # from moviepy.editor import ImageClip, concatenate_videoclips
378
+ # clips = [ImageClip(frame).set_duration(1/24) for frame in frames]
379
+ # concat_clip = concatenate_videoclips(clips, method="compose")
380
+ # concat_clip.write_videofile("test.mp4", fps=24)
381
+
382
+
383
+
384
+ img_bytes_string = turtle.save_gif("")
385
+ # turtle.save('test3.png')
386
+ with open("test5.gif", "wb") as f:
387
+ f.write(img_bytes_string)
388
+
389
+
390
+
391
+
392
+ # example_plot()
393
  # plot2()