├── .gitignore
├── README.md
├── benchmark.py
├── colorswirl.py
├── data
├── MOSHED-doctor.gif
├── arrival.json
├── frog.png
├── frog_lut.png
├── frog_src.png
├── splash.png
└── splash_output.png
├── fractal.py
├── gifreader.py
├── life.py
├── lut.py
├── pendulum.py
├── rainbowize.py
├── requirements.txt
├── screenshots
├── benchmark_results.png
├── colorswirl.gif
├── fractal_zoom.gif
├── gifreader_demo.gif
├── life_demo.gif
├── lut.gif
├── n=1000.PNG
├── n=10000.PNG
├── n=100000.PNG
├── n=1000animated.gif
├── rainbowize.gif
├── stega_demo.png
├── video_demo.gif
├── warp.gif
└── worms.gif
├── steganography.py
├── threedee.py
├── video.py
├── warp.py
└── worms.py
/.gitignore:
--------------------------------------------------------------------------------
1 | .git/*
2 | .idea/*
3 | __pycache__/*
4 | test.py
5 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # A collection of pygame experiments and utilities.
2 |
3 | All of these are free to use & modify, with or without attribution. Every top-level python file is a standalone project. Most of these require additional libraires. Use this to install every dependency, or just install the ones you need based on the imports at the top of the file you want to use.
4 | ```
5 | pip install -r requirements.txt
6 | ```
7 |
8 | # life ([life.py](life.py))
9 | An efficient game of life simulation using pygame and numpy.
10 |
11 | 
12 |
13 | After installing dependencies, run with `python life.py`
14 |
15 | # colorswirl ([colorswirl.py](colorswirl.py))
16 | A colorful celluar automata effect using pygame and numpy.
17 |
18 | 
19 |
20 | After installing dependencies, run with `python colorswirl.py`
21 |
22 | # fractal ([fractal.py](fractal.py))
23 | An implementation of the mandlebrot set, which lets you zoom in and out.
24 |
25 | 
26 |
27 | After installing dependencies, run with `python fractal.py`
28 |
29 | # rainbowize ([rainbowize.py](rainbowize.py))
30 | A function that applies a "rainbow effect" to a single surface.
31 |
32 | 
33 |
34 | After installing dependencies, run with `python rainbowize.py` to see a demo.
35 |
36 | Or import it into your own project and call `rainbowize.apply_rainbow(my_surface)`.
37 |
38 | This program also demonstrates how to set up a `pygame.SCALED` display with a custom initial scaling factor and outer fill color (see function `make_fancy_scaled_display`).
39 |
40 | # warp ([warp.py](warp.py))
41 | A function that stretches a surface into an arbitrary quad using cv2's perspective warp.
42 |
43 | 
44 |
45 | After installing dependencies (`numpy`, `cv2`), run with `python warp.py` to see a demo.
46 |
47 | # lut ([lut.py](lut.py))
48 | A function that transforms the colors of a surface using a lookup table (aka a "LUT").
49 |
50 | 
51 |
52 | After installing dependencies, run with `python lut.py` to see a demo.
53 |
54 | Or import it into your own project and call `lut.apply_lut(source_surface, lut_surface, idx)`.
55 |
56 | If `numpy` isn't available, the function will fall back to a pure pygame routine (which is slower but produces the same result). The function also has an optional built-in caching system, and handles per-pixel alpha in a reasonable way.
57 |
58 | # worms ([worms.py](worms.py))
59 | A screensaver effect made with pure pygame that creates an illusion of depth.
60 |
61 | 
62 |
63 | After installing dependencies, run with `python worms.py`
64 |
65 | # video ([video.py](video.py))
66 | A video playback utility for pygame using numpy and cv2.
67 |
68 | 
69 |
70 | See documentation in class for usage instructions.
71 |
72 | # steganography ([steganography.py](steganography.py))
73 | A module that writes text data into the pixel values of images.
74 |
75 | 
76 |
77 | After installing dependencies, import the module into your project and call its methods. See docstrings for detailed usage instructions and information about advanced settings (e.g. bit depth, image resizing).
78 |
79 | ### Example of writing a message into a surface.
80 | ```
81 | output_surface = steganography.write_text_to_surface("secret message", my_surface)
82 | message = steganography.read_text_from_surface(output_surface)
83 | print(message) # prints "secret message"
84 | ```
85 |
86 | ### Example of saving a message to a PNG.
87 | ```
88 | steganography.save_text_as_image_file("secret message", my_surface, "path/to/file.png")
89 | message = steganography.load_text_from_image_file("path/to/file.png")
90 | print(message) # prints "secret message"
91 | ```
92 |
93 | # gifreader ([gifreader.py](gifreader.py))
94 | A utility function that loads a gif file's images and metadata.
95 |
96 | Metadata includes things like frame durations and image dimensions. See documentation on the function for usage instructions.
97 |
98 | Requires `imageio` (`pip install imageio`) and `numpy`
99 |
100 | ### Example of loading and displaying a GIF
101 | 
102 |
103 | # benchmark ([benchmark.py](benchmark.py))
104 | A program that benchmarks pygame's rendering and displays the results in a graph.
105 |
106 | Requires `matplotlib` (`pip install matplotlib`).
107 |
108 | ### Instructions
109 | Use `python benchmark.py` to run the tests with default settings, or `python benchmark.py --help` to see all settings.
110 |
111 | If you close the window while the test is running, the results from cases that have already completed will still be shown.
112 |
113 | ### Sample Results
114 | 
115 |
116 | A plot that shows the relationship between FPS and the number of entities being rendered. Each line is a separate test case.
117 | ```
118 | ALL = Entities of all types rendered together
119 | SURF_RGB = Surfaces with no translucency
120 | SURF_RGBA = Surfaces with per-pixel translucency
121 | SURF_RGB_WITH_ALPHA = Surfaces with full-surface translucency
122 | RECT_FILLED = pygame.draw.rect with width = 0 (i.e. filled)
123 | CIRCLE_FILLED = pygame.draw.circle with width = 0 (i.e. filled)
124 | LINE = pygame.draw.line
125 | RECT_HOLLOW = pygame.draw.rect with width > 0
126 | CIRCLE_HOLLOW = pygame.draw.circle with width > 0
127 | ```
128 | Note that non-rectangular entities are scaled up and/or assigned widths so that drawing them will roughly change the same number of pixels as blitting a surface. This seemed like the most sensible way to compare rendering speeds.
129 |
130 | # double-pendulum ([pendulum.py](pendulum.py))
131 | An efficient double pendulum simulation using pygame, numpy, and OpenGL.
132 |
133 | ### Demo (N=1000)
134 | 
135 |
136 | ### Instructions
137 | After installing dependencies, use this to run the default program:
138 | ```
139 | python pendulum.py
140 | ```
141 |
142 | To see the program's optional arguments, use:
143 | ```
144 | python pendulum.py -h
145 | ```
146 |
147 | A command like this can be used to make a "realistic-looking" 3-pendulum simulation:
148 | ```
149 | python pendulum.py -n 3 --opacity 1.0 --size 400 400 --length 5 --zoom 20 --spread 1.3
150 | ```
151 |
152 | While the simulation is running, the following actions can be used via keybindings:
153 | * \[r\] to restart the simlution with a new starting angle
154 | * \[Esc\] to quit the program
155 | * \[p\] to enable / disable profiling
156 |
157 | ### N=1000
158 | 
159 |
160 | ### N=10,000
161 | 
162 |
163 | ### N=100,000
164 | 
165 |
--------------------------------------------------------------------------------
/benchmark.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 |
4 | import argparse
5 | import math
6 | import random
7 | import enum
8 |
9 | os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "1"
10 | import pygame
11 | import matplotlib.pyplot as plt
12 |
13 |
14 | # Pygame Benchmarking Program
15 | # by Ghast ~ github.com/davidpendergast
16 |
17 |
18 | SCREEN_SIZE = 640, 480
19 |
20 | ENT_SIZE = 64 # width & height of filled-in entities in pixels
21 |
22 | SPAWN_RATE = 200 # entities per sec, increasing this will make the tests run faster but be less accurate
23 | TICK_RATE = 0.10 # frequency of entity additions and fps measurements
24 | STOP_AT_FPS = 45 # test cases stop when this FPS is reached
25 |
26 | INCLUDE_SURFACES = True
27 | INCLUDE_FILLED_SHAPES = True
28 | INCLUDE_HOLLOW_SHAPES = True
29 | INCLUDE_LINES = True
30 |
31 | RAND_SEED = 27182818
32 |
33 | CAPTION_REFRESH_PER_SEC = 3 # per second, how frequently to update the window title
34 | LOW_FPS_GRACE_PERIOD = 0.5 # seconds, how long the FPS can stay under STOP_AT_FPS before ending the test case
35 | PAUSE_BETWEEN_TESTS_TIME = 1 # seconds, how long to pause between tests
36 | GRAPH_SMOOTHING_RADIUS = 1 # seconds, how much to smooth the graph (uses running average)
37 | LOG_AXIS = True # whether to use log-scaling for the graph's y-axis
38 |
39 | PX_PER_ENT = ENT_SIZE ** 2
40 |
41 |
42 | class EntityType(enum.IntEnum):
43 | SURF_RGB = 0 # Surface with no per-pixel alpha
44 | SURF_RGB_WITH_ALPHA = 1 # Surface with no per-pixel alpha, but an alpha value < 255
45 | SURF_RGBA = 2 # Surface with per-pixel alpha.
46 |
47 | LINE = 3 # A line drawn via pygame.draw.line
48 | RECT_FILLED = 4 # A rect drawn via pygame.draw.rect (with width = 0)
49 | RECT_HOLLOW = 5 # A rect drawn via pygame.draw.rect (with width > 0)
50 | CIRCLE_FILLED = 6 # A circle drawn via pygame.draw.circle (with width = 0)
51 | CIRCLE_HOLLOW = 7 # A circle drawn via pygame.draw.circle (with width > 0)
52 |
53 |
54 | def _calc_avg_lengths(w, h, n=10000):
55 | """Finds the average length of a line between two random points in a (w x h) area using random sampling.
56 | returns: Average total length, x-length, and y-length
57 | """
58 | total_dx = 0
59 | total_dy = 0
60 | total_dist = 0
61 |
62 | # Fixed seed because this can affect the test quite a lot. For example, if hollow circles
63 | # use a thickness of 6 pixels instead of 5 that's a 20% increase.
64 | rand = random.Random(x=12345)
65 |
66 | for _ in range(n):
67 | p1 = rand.randint(0, w - 1), rand.randint(0, h - 1)
68 | p2 = rand.randint(0, w - 1), rand.randint(0, h - 1)
69 | dx = p1[0] - p2[0]
70 | dy = p1[1] - p2[1]
71 | total_dx += abs(dx)
72 | total_dy += abs(dy)
73 | total_dist += math.sqrt(dx * dx + dy * dy)
74 | return total_dist / n, total_dx / n, total_dy / n
75 |
76 |
77 | AVG_LENGTH, AVG_WIDTH, AVG_HEIGHT = _calc_avg_lengths(*SCREEN_SIZE)
78 |
79 | # Want to define the geometric entities such that on average they'll have an area of about PX_PER_ENT pixels.
80 | # Note that these are approximate (particularly the circle one, which assumes the relationship between radius
81 | # and pixels changed is linear (it's actually quadratic)).
82 | LINE_WIDTH = max(1, round(PX_PER_ENT / AVG_LENGTH))
83 | HOLLOW_RECT_WIDTH = max(1, round(PX_PER_ENT / (2 * (AVG_WIDTH + AVG_HEIGHT))))
84 | HOLLOW_CIRCLE_RADIUS_MULT = 0.25 # radius = mult * dist between two random points
85 | HOLLOW_CIRCLE_WIDTH = max(1, round((AVG_LENGTH * HOLLOW_CIRCLE_RADIUS_MULT)
86 | - math.sqrt(max(0, (AVG_LENGTH * HOLLOW_CIRCLE_RADIUS_MULT) ** 2
87 | - PX_PER_ENT / math.pi))))
88 | FILLED_CIRCLE_RADIUS = max(1, round(math.sqrt(PX_PER_ENT / math.pi)))
89 |
90 |
91 | def _print_info():
92 | print(f"\nStarting new simulation:\n"
93 | f" Screen size = {SCREEN_SIZE} px\n"
94 | f" Entity size = {ENT_SIZE} x {ENT_SIZE} px\n"
95 | f" Minimum FPS = {STOP_AT_FPS} fps\n"
96 | f" Spawn Rate = {SPAWN_RATE} ents/sec\n"
97 | f" Tick Rate = {TICK_RATE} sec")
98 |
99 | print(f"\nAverage number of pixels changed per render:")
100 | print(f" Filled Rect: {PX_PER_ENT:.2f} = {ENT_SIZE}**2")
101 | print(f" Filled Circle: {math.pi * FILLED_CIRCLE_RADIUS**2:.2f} = pi * {FILLED_CIRCLE_RADIUS} ** 2")
102 | print(f" Line: {AVG_LENGTH * LINE_WIDTH:.2f} = {AVG_LENGTH:.2f} * {LINE_WIDTH:.2f}")
103 | print(f" Hollow Rect: {2 * (AVG_WIDTH + AVG_HEIGHT) * HOLLOW_RECT_WIDTH:.2f}"
104 | f" = 2 * ({AVG_WIDTH:.2f} + {AVG_HEIGHT:.2f}) * {HOLLOW_RECT_WIDTH:.2f}")
105 | avg_hollow_circle_area = math.pi * ((AVG_LENGTH * HOLLOW_CIRCLE_RADIUS_MULT)**2
106 | - (AVG_LENGTH*HOLLOW_CIRCLE_RADIUS_MULT - HOLLOW_CIRCLE_WIDTH)**2)
107 | print(f" Hollow Circle: {avg_hollow_circle_area:.2f}"
108 | f" = pi * (({AVG_LENGTH:.2f})**2 - ({AVG_LENGTH:.2f}"
109 | f" - {HOLLOW_CIRCLE_WIDTH:.2f})**2) (approx.)")
110 |
111 |
112 | class EntityFactory:
113 |
114 | def __init__(self, seed=RAND_SEED):
115 | self.rand = random.Random(x=seed)
116 |
117 | def get_next(self):
118 | return (random.randint(0, 4096), # determines Entity type
119 | random.randint(0, 4096), # determines (x1, y1)
120 | random.randint(0, 4096), # determines (x2, y2)
121 | random.randint(0, 4096)) # determines color and/or opacity
122 |
123 |
124 | class Renderer:
125 |
126 | def __init__(self, screen, n_pts=503, ent_types=tuple(e for e in EntityType), seed=RAND_SEED):
127 | self.screen = screen
128 | self.entities = []
129 | self.ent_types = ent_types
130 | self.random = random.Random(x=seed)
131 |
132 | # Instantiate the invisible points that float around the screen.
133 | # These points determine the locations of entities. Note that they're not
134 | # updated one-by-one each frame -- instead, their positions are derived
135 | # on-the-fly using their initial position, velocity, and the time.
136 | w, h = self.screen.get_size()
137 | self.pts = [pygame.Vector2(self.random.randint(0, w-1),
138 | self.random.randint(0, h-1)) for _ in range(n_pts)]
139 | self.vels = [pygame.Vector2(0, 1) for _ in range(n_pts)]
140 | for v in self.vels:
141 | v.rotate_ip(random.random() * 360.0)
142 | v.scale_to_length(random.randint(30, 50))
143 |
144 | self.t = 0 # current time
145 | self.colors = ["red", "green", "blue", "yellow", "cyan", "magenta", "orange"]
146 |
147 | # Precompute the surfaces that entities may need.
148 | self.rgb_surfs = []
149 | self.rgba_surfs = []
150 | self.rgb_surfs_with_alpha = []
151 | for idx, c in enumerate(self.colors):
152 | rgb_surf = pygame.Surface((ENT_SIZE, ENT_SIZE))
153 | rgb_surf.fill(c)
154 | self.rgb_surfs.append(rgb_surf)
155 |
156 | opacity = (idx + 1) / len(self.colors)
157 |
158 | rgba_surf = pygame.Surface((ENT_SIZE, ENT_SIZE), flags=pygame.SRCALPHA)
159 | rgb = tuple(i for i in pygame.color.Color(c))
160 | rgba_surf.fill((*rgb[:3], int(opacity * 255)))
161 | self.rgba_surfs.append(rgba_surf)
162 |
163 | rgb_surf_with_alpha = pygame.Surface((ENT_SIZE, ENT_SIZE))
164 | rgb_surf_with_alpha.fill(c)
165 | rgb_surf_with_alpha.set_alpha(int(opacity * 255))
166 | self.rgb_surfs_with_alpha.append(rgb_surf_with_alpha)
167 |
168 | # Bit hacky, but it's important to look up each entity's render method quickly.
169 | self.render_methods = [None] * len(EntityType)
170 | for e in EntityType:
171 | self.render_methods[e] = getattr(self, f"render_{e.name}")
172 |
173 | def update(self, dt):
174 | self.t += dt
175 |
176 | def get_point(self, n1, n2):
177 | base_pt = self.pts[n1 % len(self.pts)] + self.pts[n2 % len(self.pts)]
178 | vel = self.vels[n2 % len(self.vels)] + self.vels[n1 % len(self.vels)]
179 | pt = base_pt + vel * self.t
180 | x = int(pt.x) % self.screen.get_width()
181 | y = int(pt.y) % self.screen.get_height()
182 | return x, y
183 |
184 | def render(self, bg_color=(0, 0, 0)):
185 | self.screen.fill(bg_color)
186 | for ent in self.entities:
187 | ent_type = self.ent_types[ent[0] % len(self.ent_types)]
188 | p1 = self.get_point(ent[0], ent[1])
189 | p2 = self.get_point(ent[0], ent[2])
190 | color_idx = ent[3]
191 |
192 | self.render_methods[ent_type](p1, p2, color_idx) # do actual rendering
193 |
194 | def render_SURF_RGB(self, p1, p2, color_idx):
195 | surf = self.rgb_surfs[color_idx % len(self.rgb_surfs)]
196 | self.screen.blit(surf, (p1[0] - ENT_SIZE // 2, p1[1] - ENT_SIZE // 2))
197 |
198 | def render_SURF_RGBA(self, p1, p2, color_idx):
199 | surf = self.rgba_surfs[color_idx % len(self.rgba_surfs)]
200 | self.screen.blit(surf, (p1[0] - ENT_SIZE // 2, p1[1] - ENT_SIZE // 2))
201 |
202 | def render_SURF_RGB_WITH_ALPHA(self, p1, p2, color_idx):
203 | surf = self.rgb_surfs_with_alpha[color_idx % len(self.rgb_surfs_with_alpha)]
204 | self.screen.blit(surf, (p1[0] - ENT_SIZE // 2, p1[1] - ENT_SIZE // 2))
205 |
206 | def render_LINE(self, p1, p2, color_idx):
207 | c = self.colors[color_idx % len(self.colors)]
208 | pygame.draw.line(self.screen, c, p1, p2, width=LINE_WIDTH)
209 |
210 | def render_RECT_HOLLOW(self, p1, p2, color_idx):
211 | c = self.colors[color_idx % len(self.colors)]
212 | pygame.draw.rect(self.screen, c, (p1[0], p1[1], p2[0] - p1[0], p2[1] - p1[1]),
213 | width=HOLLOW_RECT_WIDTH)
214 |
215 | def render_RECT_FILLED(self, p1, p2, color_idx):
216 | c = self.colors[color_idx % len(self.colors)]
217 | pygame.draw.rect(self.screen, c, (p1[0] - ENT_SIZE // 2, p1[1] - ENT_SIZE // 2, ENT_SIZE, ENT_SIZE))
218 |
219 | def render_CIRCLE_HOLLOW(self, p1, p2, color_idx):
220 | c = self.colors[color_idx % len(self.colors)]
221 | r = (abs(p2[0] - p1[0]) + abs(p2[1] - p1[0])) * HOLLOW_CIRCLE_RADIUS_MULT
222 | pygame.draw.circle(self.screen, c, p1, r, width=HOLLOW_CIRCLE_WIDTH)
223 |
224 | def render_CIRCLE_FILLED(self, p1, p2, color_idx):
225 | c = self.colors[color_idx % len(self.colors)]
226 | pygame.draw.circle(self.screen, c, p1, FILLED_CIRCLE_RADIUS)
227 |
228 |
229 | def start_plot(title, subtitle=None):
230 | if subtitle is not None:
231 | plt.suptitle(title, fontsize=16)
232 | plt.title(subtitle, fontsize=10, y=1)
233 | else:
234 | plt.title(title)
235 | plt.xlabel('Entities')
236 | plt.ylabel('FPS')
237 | if LOG_AXIS:
238 | plt.yscale('log')
239 | yticks = [15, 30, 45, 60, 120, 144, 240]
240 | plt.yticks(yticks, [str(yt) for yt in yticks])
241 |
242 |
243 | def get_x_and_y(data, smooth_radius=0):
244 | x = []
245 | y = []
246 | for cnt in data:
247 | x.append(cnt)
248 | y.append(data[cnt])
249 | if smooth_radius > 0 and len(x) > 0:
250 | y = smooth_data(x, y, smooth_radius)
251 | return x, y
252 |
253 |
254 | def add_to_plot(data, label, show_t60=False, smooth_radius=0):
255 | x, y = get_x_and_y(data, smooth_radius=smooth_radius)
256 |
257 | if show_t60:
258 | t60 = solve_for_t(x, y, 60)
259 | if t60 is not None:
260 | plt.axhline(y=60, xmin=0, xmax=t60, color='red', linestyle='dotted', linewidth=1)
261 | plt.axvline(x=t60, color='red', linestyle='dotted', linewidth=1)
262 |
263 | if "SURF" in label:
264 | linestyle = "solid"
265 | if "ALPHA" in label or "RGBA" in label:
266 | linewidth = 1.5
267 | else:
268 | linewidth = 2
269 |
270 | elif "CIRCLE" in label or "RECT" in label or "LINE" in label:
271 | linestyle = "dashed"
272 | if "HOLLOW" in label:
273 | linewidth = 1
274 | else:
275 | linewidth = 1.5
276 | else:
277 | linestyle = "dashdot"
278 | linewidth = 2
279 |
280 | plt.plot(x, y, label=label, linestyle=linestyle, linewidth=linewidth)
281 |
282 |
283 | def finish_plot():
284 | plt.legend()
285 | plt.get_current_fig_manager().set_window_title('Benchmark Results')
286 | plt.show()
287 |
288 |
289 | def solve_for_t(x, y, target_y, or_else=None):
290 | res = None
291 | for i in range(0, len(x) - 1):
292 | if y[i] > target_y >= y[i + 1]:
293 | a = (target_y - y[i + 1]) / (y[i] - y[i + 1])
294 | res = x[i + 1] + a * (x[i] - x[i + 1])
295 | break
296 | return res if res is not None else or_else
297 |
298 |
299 | def smooth_data(x, y, box_radius):
300 | y_smoothed = []
301 | for i in range(len(x)):
302 | y_sum, n_pts = y[i], 1
303 | bs = min([box_radius, abs(x[i] - x[0]), abs(x[-1] - x[i])])
304 |
305 | i2 = i - 1
306 | while i2 >= 0 and abs(x[i2] - x[i]) <= bs:
307 | y_sum += y[i2]
308 | n_pts += 1
309 | i2 -= 1
310 |
311 | i2 = i + 1
312 | while i2 < len(x) and abs(x[i2] - x[i]) <= bs:
313 | y_sum += y[i2]
314 | n_pts += 1
315 | i2 += 1
316 |
317 | y_smoothed.append(y_sum / n_pts)
318 | return y_smoothed
319 |
320 |
321 | class TestCase:
322 |
323 | def __init__(self, name, caption_title, screen, ent_types=tuple(e for e in EntityType), seed=RAND_SEED):
324 | self.name = name
325 | self.caption_title = caption_title
326 | self.screen = screen
327 | self.ent_types = ent_types
328 |
329 | self.seed = seed
330 | self.factory = None
331 | self.renderer = None
332 | self.clock = None
333 |
334 | def start(self, pause=2):
335 | pygame.display.set_caption(f"{self.caption_title}".replace("#", "N=0"))
336 | self.factory = EntityFactory(seed=self.seed)
337 | self.renderer = Renderer(self.screen, ent_types=self.ent_types, seed=self.seed)
338 | self.renderer.t = -pause
339 |
340 | self.clock = pygame.time.Clock()
341 | last_t_above_stopping_point = 0
342 | dt = 0
343 |
344 | results = {}
345 |
346 | running = True
347 | while running:
348 | for evt in pygame.event.get():
349 | if evt.type == pygame.QUIT:
350 | raise ValueError("Quit")
351 |
352 | self.renderer.update(dt)
353 | self.renderer.render()
354 | pygame.display.flip()
355 |
356 | dt = self.clock.tick() / 1000.0
357 |
358 | n = max(0, int(int(self.renderer.t / TICK_RATE) * TICK_RATE * SPAWN_RATE))
359 | if len(self.renderer.entities) < n:
360 | last_n = len(self.renderer.entities)
361 | if last_n > 0:
362 | results[last_n] = self.clock.get_fps()
363 | if results[last_n] >= STOP_AT_FPS:
364 | last_t_above_stopping_point = self.renderer.t
365 | elif self.renderer.t > last_t_above_stopping_point + LOW_FPS_GRACE_PERIOD:
366 | running = False
367 | while len(self.renderer.entities) < n:
368 | self.renderer.entities.append(self.factory.get_next())
369 |
370 | # update caption periodically
371 | cap_refresh = CAPTION_REFRESH_PER_SEC
372 | if int((self.renderer.t + dt) * cap_refresh) > int(self.renderer.t * cap_refresh):
373 | pygame.display.set_caption(self.caption_title.replace("#", f"N={n}, FPS={self.clock.get_fps():.2f}"))
374 |
375 | return results
376 |
377 |
378 | def _print_result(case_num, test, res, fps_to_display=(144, 120, 60, STOP_AT_FPS)):
379 | print(f"\n{case_num}. {test.name} Results:")
380 | x, y = get_x_and_y(res, smooth_radius=GRAPH_SMOOTHING_RADIUS * SPAWN_RATE)
381 | for fps in reversed(sorted(f for f in set(fps_to_display) if f >= STOP_AT_FPS)):
382 | t = solve_for_t(x, y, fps, or_else=0)
383 | print(f" {fps:>3} FPS: {round(t):>4} entities")
384 |
385 |
386 | def _build_test_cases(screen):
387 | to_exclude = set()
388 | if not INCLUDE_SURFACES:
389 | to_exclude.update([EntityType.SURF_RGB, EntityType.SURF_RGBA, EntityType.SURF_RGB_WITH_ALPHA])
390 | if not INCLUDE_FILLED_SHAPES:
391 | to_exclude.update([EntityType.RECT_FILLED, EntityType.CIRCLE_FILLED])
392 | if not INCLUDE_LINES:
393 | to_exclude.add(EntityType.LINE)
394 | if not INCLUDE_HOLLOW_SHAPES:
395 | to_exclude.update([EntityType.RECT_HOLLOW, EntityType.CIRCLE_HOLLOW])
396 |
397 | ents_to_test = [e for e in EntityType if e not in to_exclude]
398 |
399 | test_cases = []
400 | if len(ents_to_test) == 0:
401 | raise ValueError("Nothing to test, all cases are disabled.")
402 | elif len(ents_to_test) == 1:
403 | e = ents_to_test[0]
404 | test_cases.append(TestCase(e.name, f"{e.name} (#, CASE=1/1)", screen, (e,)))
405 | else:
406 | n_cases = 1 + len(ents_to_test)
407 | test_cases = [TestCase("ALL", f"ALL (#, CASE=1/{n_cases})", screen, ents_to_test)]
408 | for e_idx, e in enumerate(ents_to_test):
409 | test_cases.append(TestCase(e.name, f"{e.name} (#, CASE={2 + e_idx}/{n_cases})", screen, (e,)))
410 | return test_cases
411 |
412 |
413 | def _run():
414 | pygame.init()
415 | screen = pygame.display.set_mode(SCREEN_SIZE)
416 |
417 | _print_info()
418 | test_cases = _build_test_cases(screen)
419 |
420 | all_results = {}
421 | try:
422 | for test in test_cases:
423 | res = test.start(pause=PAUSE_BETWEEN_TESTS_TIME)
424 | all_results[test.name] = res
425 | _print_result(test_cases.index(test) + 1, test, res)
426 | except ValueError as err:
427 | if str(err) == "Quit":
428 | msg = "no results to show" if len(all_results) == 0 else "showing partial results"
429 | print(f"\nBenchmark was cancelled before completion ({msg})")
430 | pass # show partial results (if possible) when you quit
431 | else:
432 | raise err
433 |
434 | pygame.display.quit()
435 |
436 | # display the plot
437 | if len(all_results) > 0:
438 | print("\nDisplaying plot...")
439 | pg_ver = pygame.version.ver
440 | sdl_ver = ".".join(str(v) for v in pygame.version.SDL)
441 | py_ver = ".".join(str(v) for v in sys.version_info[:3])
442 | start_plot(f"FPS vs. Entities ({ENT_SIZE}x{ENT_SIZE})",
443 | subtitle=f"pygame {pg_ver} (SDL {sdl_ver}, Python {py_ver})")
444 | for test_name in all_results:
445 | add_to_plot(all_results[test_name], test_name, show_t60=(test_name == "ALL"),
446 | smooth_radius=GRAPH_SMOOTHING_RADIUS * SPAWN_RATE)
447 | finish_plot()
448 |
449 |
450 | if __name__ == "__main__":
451 | parser = argparse.ArgumentParser(description='A benchmarking program for pygame rendering.')
452 | parser.add_argument('--size', type=int, metavar="int", default=SCREEN_SIZE, nargs=2, help=f'the window size (default {SCREEN_SIZE[0]} {SCREEN_SIZE[1]})')
453 | parser.add_argument('--entity-size', type=int, metavar="int", default=ENT_SIZE, help=f'the size of entities in both dimensions (default {ENT_SIZE} px)')
454 |
455 | parser.add_argument('--skip-surfaces', dest='surfaces', action='store_false', default=True, help=f'if used, will skip tests for RGB, RGBA, and RGB (with alpha) surfaces')
456 | parser.add_argument('--skip-filled', dest='filled', action='store_false', default=True, help=f'if used, will skip tests for pygame.draw.rect and circle with width = 0')
457 | parser.add_argument('--skip-hollow', dest='hollow', action='store_false', default=True, help=f'if used, will skip tests for pygame.draw.rect and circle with width > 0')
458 | parser.add_argument('--skip-lines', dest='lines', action='store_false', default=True, help=f'if used, will skip tests for pygame.draw.line')
459 |
460 | parser.add_argument('--spawn-rate', type=int, metavar="int", default=SPAWN_RATE, help=f'number of entities to spawn per second (default {SPAWN_RATE}, smaller = slower and more accurate)')
461 | parser.add_argument('--tick-rate', type=float, metavar="float", default=TICK_RATE, help=f'how frequently to sample the FPS and add new entities (in seconds, default {TICK_RATE})')
462 | parser.add_argument('--fps-thresh', type=int, metavar="int", default=STOP_AT_FPS, help=f'the FPS at which a test case should stop (default {STOP_AT_FPS})')
463 |
464 | parser.add_argument('--smooth', type=float, metavar="float", default=GRAPH_SMOOTHING_RADIUS, help=f'how much to smooth the graph, in seconds (default {GRAPH_SMOOTHING_RADIUS}, or use 0 for none)')
465 | parser.add_argument('--no-log-axis', dest='log_axis', action='store_false', default=LOG_AXIS, help=f'if used, will disable log-scaling on the graph\'s y-axis')
466 |
467 | parser.add_argument('--seed', type=int, metavar="int", default=RAND_SEED, help=f'random seed (default={RAND_SEED}, use 0 to generate one)')
468 |
469 | args = parser.parse_args()
470 |
471 | SCREEN_SIZE = args.size
472 | ENT_SIZE = args.entity_size
473 |
474 | INCLUDE_SURFACES = args.surfaces
475 | INCLUDE_FILLED_SHAPES = args.filled
476 | INCLUDE_HOLLOW_SHAPES = args.hollow
477 | INCLUDE_LINES = args.lines
478 |
479 | SPAWN_RATE = args.spawn_rate
480 | TICK_RATE = args.tick_rate
481 | STOP_AT_FPS = args.fps_thresh
482 | GRAPH_SMOOTHING_RADIUS = args.smooth
483 | LOG_AXIS = args.log_axis
484 | RAND_SEED = args.seed if args.seed > 0 else None
485 |
486 | _run()
487 |
--------------------------------------------------------------------------------
/colorswirl.py:
--------------------------------------------------------------------------------
1 | import pygame
2 | import numpy
3 |
4 | W, H = 600, 300
5 | FPS = 60
6 |
7 | FLAGS = 0 # | pygame.SCALED
8 |
9 |
10 | class State:
11 | def __init__(self):
12 | self.grid = numpy.random.randint(0, 0xFFFFFF + 1, (W, H), numpy.int32)
13 |
14 | def step(self):
15 | shifted = []
16 | for dx in [-1, 0, 1]:
17 | for dy in [-1, 0, 1]:
18 | if (dx, dy) != (0, 0):
19 | shifted.append(numpy.roll(self.grid, (dx, dy), (0, 1)))
20 |
21 | rand_vals = numpy.random.randint(0, len(shifted), (W, H), numpy.int16)
22 |
23 | masks = []
24 | for i in range(len(shifted)):
25 | masks.append(1 * (rand_vals == i))
26 |
27 | self.grid[:] = 0
28 | for i in range(len(shifted)):
29 | self.grid += masks[i] * shifted[i]
30 |
31 | def draw(self, screen):
32 | pygame.surfarray.blit_array(screen, self.grid)
33 |
34 |
35 | if __name__ == "__main__":
36 | pygame.init()
37 | screen = pygame.display.set_mode((W, H), flags=FLAGS)
38 | state = State()
39 | clock = pygame.time.Clock()
40 |
41 | running = True
42 | last_update_time = pygame.time.get_ticks()
43 |
44 | while running:
45 | current_time = pygame.time.get_ticks()
46 | dt = (current_time - last_update_time) / 1000
47 | for e in pygame.event.get():
48 | if e.type == pygame.QUIT or (e.type == pygame.KEYDOWN and e.key == pygame.K_ESCAPE):
49 | running = False
50 | elif e.type == pygame.KEYDOWN:
51 | if e.key == pygame.K_r:
52 | print("Resetting.")
53 | state = State()
54 | elif e.key == pygame.K_RIGHT:
55 | FPS += 10
56 | print("Increased target FPS to: {}".format(FPS))
57 | elif e.key == pygame.K_LEFT:
58 | FPS = max(10, FPS - 10)
59 | print("Decreased target FPS to: {}".format(FPS))
60 | state.step()
61 | state.draw(screen)
62 |
63 | pygame.display.flip()
64 |
65 | if current_time // 1000 > last_update_time // 1000:
66 | pygame.display.set_caption("Color Swirl (FPS={:.1f}, TARGET_FPS={}, SIZE={})".format(clock.get_fps(), FPS, (W, H)))
67 |
68 | last_update_time = current_time
69 | clock.tick(FPS)
70 |
--------------------------------------------------------------------------------
/data/MOSHED-doctor.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/data/MOSHED-doctor.gif
--------------------------------------------------------------------------------
/data/arrival.json:
--------------------------------------------------------------------------------
1 | {
2 | "characters": ["player_fast", "player_small"],
3 | "description": "The journey begins.",
4 | "entities": [
5 | {"art_id": 0, "color_id": 1, "h": 16, "type": "block", "w": 64, "x": 32, "y": 80},
6 | {"color_id": 1, "h": 80, "type": "block", "w": 16, "x": 0, "y": 48},
7 | {"color_id": 2, "h": 16, "type": "block", "w": 64, "x": 64, "y": 224},
8 | {"color_id": 1, "h": 32, "type": "block", "w": 80, "x": 16, "y": 96},
9 | {"h": 16, "subtype": "player_fast", "type": "start_block", "w": 16, "x": 48, "x_dir": 1, "y": 224},
10 | {"h": 16, "subtype": "player_small", "type": "start_block", "w": 16, "x": 16, "x_dir": 1, "y": 80},
11 | {"art_id": 0, "color_id": 1, "h": 48, "type": "block", "w": 32, "x": 224, "y": 176},
12 | {"color_id": 1, "h": 48, "type": "block", "w": 80, "x": 0, "y": 0},
13 | {"art_id": 1, "color_id": 1, "h": 16, "type": "block", "w": 32, "x": 288, "y": 224},
14 | {"art_id": 1, "color_id": 2, "h": 16, "type": "block", "w": 32, "x": 16, "y": 224},
15 | {"art_id": 1, "color_id": 2, "h": 32, "type": "block", "w": 16, "x": 16, "y": 192},
16 | {
17 | "art_id": 2,
18 | "color_id": 1,
19 | "h": 32,
20 | "subtype": [0, 0, "horizontal"],
21 | "type": "sloped_2x2_block",
22 | "w": 32,
23 | "x": 96,
24 | "y": 80
25 | },
26 | {
27 | "art_id": 1,
28 | "color_id": 2,
29 | "h": 32,
30 | "subtype": [2, 0, "horizontal"],
31 | "type": "sloped_2x2_block",
32 | "w": 32,
33 | "x": 192,
34 | "y": 192
35 | },
36 | {
37 | "art_id": 0,
38 | "color_id": 2,
39 | "h": 32,
40 | "subtype": [2, 0, "horizontal"],
41 | "type": "sloped_2x2_block",
42 | "w": 32,
43 | "x": 160,
44 | "y": 208
45 | },
46 | {"art_id": 0, "color_id": 1, "h": 16, "type": "block", "w": 64, "x": 192, "y": 224},
47 | {"art_id": 0, "color_id": 1, "h": 32, "type": "block", "w": 64, "x": 416, "y": 208},
48 | {
49 | "art_id": 1,
50 | "color_id": 1,
51 | "h": 32,
52 | "subtype": [0, 0, "horizontal"],
53 | "type": "sloped_2x2_block",
54 | "w": 32,
55 | "x": 256,
56 | "y": 208
57 | },
58 | {"art_id": 0, "color_id": 1, "h": 16, "type": "block", "w": 48, "x": 80, "y": 128},
59 | {"art_id": 0, "color_id": 1, "h": 16, "type": "block", "w": 8, "x": 120, "y": 144},
60 | {"art_id": 0, "color_id": 1, "h": 8, "type": "block", "w": 8, "x": 112, "y": 144},
61 | {"art_id": 1, "color_id": 1, "h": 32, "type": "block", "w": 32, "x": 224, "y": 48},
62 | {"art_id": 0, "color_id": 1, "h": 8, "type": "block", "w": 16, "x": 240, "y": 80},
63 | {"art_id": 0, "color_id": 1, "h": 16, "type": "block", "w": 16, "x": 208, "y": 48},
64 | {"art_id": 0, "color_id": 1, "h": 8, "type": "block", "w": 16, "x": 272, "y": 48},
65 | {"art_id": 0, "color_id": 1, "h": 8, "type": "block", "w": 48, "x": 416, "y": 48},
66 | {"art_id": 0, "color_id": 1, "h": 16, "type": "block", "w": 16, "x": 464, "y": 48},
67 | {"art_id": 0, "h": 16, "subtype": 1, "type": "door_block", "w": 32, "x": 256, "y": 176},
68 | {"subtype": 1, "type": "key", "x": 176, "y": 128},
69 | {"art_id": 0, "color_id": 1, "h": 16, "type": "block", "w": 16, "x": 256, "y": 48},
70 | {"art_id": 0, "color_id": 2, "h": 16, "type": "block", "w": 32, "x": 128, "y": 224},
71 | {"color_id": -1, "h": 16, "subtype": "player_small", "type": "end_block", "w": 32, "x": 432, "y": 192},
72 | {
73 | "art_id": 2,
74 | "color_id": 2,
75 | "h": 32,
76 | "subtype": [2, 0, "horizontal"],
77 | "type": "sloped_2x2_block",
78 | "w": 32,
79 | "x": 288,
80 | "y": 160
81 | },
82 | {"art_id": 0, "color_id": 2, "h": 144, "type": "block", "w": 16, "x": 464, "y": 64},
83 | {
84 | "art_id": 2,
85 | "color_id": 1,
86 | "h": 32,
87 | "subtype": [2, 0, "horizontal"],
88 | "type": "sloped_2x2_block",
89 | "w": 32,
90 | "x": 384,
91 | "y": 208
92 | },
93 | {"art_id": 0, "color_id": 1, "h": 16, "type": "block", "w": 64, "x": 320, "y": 224},
94 | {
95 | "art_id": 1,
96 | "color_id": 2,
97 | "h": 32,
98 | "subtype": [0, 2, "horizontal"],
99 | "type": "sloped_2x2_block",
100 | "w": 32,
101 | "x": 352,
102 | "y": 160
103 | },
104 | {
105 | "art_id": 0,
106 | "color_id": 2,
107 | "h": 32,
108 | "subtype": [2, 0, "horizontal"],
109 | "type": "sloped_2x2_block",
110 | "w": 32,
111 | "x": 352,
112 | "y": 128
113 | },
114 | {
115 | "art_id": 0,
116 | "color_id": 2,
117 | "h": 32,
118 | "subtype": [0, 2, "horizontal"],
119 | "type": "sloped_2x2_block",
120 | "w": 32,
121 | "x": 384,
122 | "y": 144
123 | },
124 | {
125 | "art_id": 0,
126 | "color_id": 2,
127 | "h": 32,
128 | "subtype": [2, 0, "horizontal"],
129 | "type": "sloped_2x2_block",
130 | "w": 32,
131 | "x": 384,
132 | "y": 112
133 | },
134 | {
135 | "art_id": 0,
136 | "color_id": 2,
137 | "h": 32,
138 | "subtype": [2, 0, "horizontal"],
139 | "type": "sloped_2x2_block",
140 | "w": 32,
141 | "x": 384,
142 | "y": 112
143 | },
144 | {"art_id": 2, "color_id": 2, "h": 16, "type": "block", "w": 16, "x": 416, "y": 112},
145 | {
146 | "art_id": 0,
147 | "color_id": 2,
148 | "h": 32,
149 | "subtype": [2, 0, "horizontal"],
150 | "type": "sloped_2x2_block",
151 | "w": 32,
152 | "x": 320,
153 | "y": 144
154 | },
155 | {"art_id": 0, "color_id": 2, "h": 16, "type": "block", "w": 32, "x": 320, "y": 176},
156 | {"color_id": -1, "h": 16, "subtype": "player_fast", "type": "end_block", "w": 32, "x": 432, "y": 112},
157 | {
158 | "art_id": 0,
159 | "color_id": 2,
160 | "h": 32,
161 | "subtype": [0, 2, "horizontal"],
162 | "type": "sloped_2x2_block",
163 | "w": 32,
164 | "x": 416,
165 | "y": 128
166 | },
167 | {"art_id": 2, "color_id": 2, "h": 16, "type": "block", "w": 16, "x": 448, "y": 128},
168 | {"color_id": 1, "h": 48, "type": "block", "w": 80, "x": 400, "y": 0},
169 | {"color_id": 1, "h": 48, "type": "block", "w": 320, "x": 80, "y": 0},
170 | {"art_id": 0, "color_id": 1, "h": 48, "type": "block", "w": 16, "x": 0, "y": 192},
171 | {"art_id": 1, "color_id": 1, "h": 16, "type": "block", "w": 32, "x": 96, "y": 112},
172 | {"art_id": 1, "color_id": 1, "h": 16, "type": "block", "w": 16, "x": 128, "y": 112},
173 | {"art_id": 2, "color_id": 1, "h": 16, "type": "block", "w": 16, "x": 176, "y": 144},
174 | {"color_id": 1, "h": 64, "type": "block", "w": 16, "x": 0, "y": 128}
175 | ],
176 | "level_id": "arrival",
177 | "name": "Arrival",
178 | "song_id": null,
179 | "special": ["show_instructions"],
180 | "time_limit": 3600
181 | }
--------------------------------------------------------------------------------
/data/frog.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/data/frog.png
--------------------------------------------------------------------------------
/data/frog_lut.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/data/frog_lut.png
--------------------------------------------------------------------------------
/data/frog_src.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/data/frog_src.png
--------------------------------------------------------------------------------
/data/splash.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/data/splash.png
--------------------------------------------------------------------------------
/data/splash_output.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/data/splash_output.png
--------------------------------------------------------------------------------
/fractal.py:
--------------------------------------------------------------------------------
1 |
2 | import pygame
3 | import numpy
4 |
5 | W, H = 600, 450
6 |
7 | FPS = 60
8 | TOTAL_ITERATIONS_PER_FRAME = 100_000
9 | MAX_ITERS_PER_CELL_PER_FRAME = 250
10 |
11 | SHOW_BOX = True # B to toggle
12 | SIMPLE_COLORING = False # C to toggle
13 |
14 | FLAGS = 0 # | pygame.SCALED
15 |
16 | ABS_LIMIT = 2
17 |
18 |
19 | class State:
20 |
21 | def __init__(self, view_rect=(-2, -2 * (H / W), 4, 4 * (H / W))):
22 | self.colors = numpy.zeros((W, H), numpy.int32)
23 |
24 | self.rect = view_rect
25 | self.iter_counts = numpy.zeros((W, H))
26 |
27 | self.last_iters_0 = 1
28 | self.last_iters_1 = 0
29 |
30 | self.min_escaped_iter = 0
31 | self.max_escaped_iter = 1
32 |
33 | yy, xx = numpy.meshgrid([self.rect[1] + i * self.rect[3] / H for i in range(0, H)],
34 | [self.rect[0] + i * self.rect[2] / W for i in range(0, W)])
35 |
36 | self.c = xx - 1j * yy
37 | self.z = numpy.zeros((W, H), dtype=complex)
38 |
39 | self.cells_to_compute = numpy.abs(self.z) < ABS_LIMIT
40 |
41 | def update(self, total_iterations=TOTAL_ITERATIONS_PER_FRAME):
42 | cells = self.cells_to_compute
43 | n_cells = numpy.count_nonzero(cells)
44 |
45 | if n_cells > 0:
46 | iters_per_cell = min(max(1, int(total_iterations / n_cells)), MAX_ITERS_PER_CELL_PER_FRAME)
47 | for i in range(iters_per_cell):
48 | self.step(cells)
49 |
50 | if numpy.any(~cells):
51 | self.max_escaped_iter = numpy.max(self.iter_counts[~cells])
52 | self.min_escaped_iter = numpy.min(self.iter_counts[~cells])
53 | if self.max_escaped_iter == self.min_escaped_iter:
54 | self.max_escaped_iter += 1
55 |
56 | self.update_colors()
57 |
58 | def get_deepest_iteration(self):
59 | return int(numpy.max(self.iter_counts))
60 |
61 | def get_pretty_rect(self) -> str:
62 | return "[{:.2f}, {:.2f}, {:.4f}, {:.4f}]".format(*self.rect)
63 |
64 | def step(self, cells):
65 | self.z[cells] *= self.z[cells]
66 | self.z[cells] += self.c[cells]
67 | self.iter_counts[cells] += 1
68 |
69 | cells[cells] = numpy.abs(self.z[cells]) < ABS_LIMIT
70 |
71 | def update_colors(self):
72 | map_to_rainbow(self.iter_counts, self.colors, self.min_escaped_iter, self.max_escaped_iter)
73 |
74 | in_set = self.cells_to_compute
75 | self.colors[in_set] = 0
76 |
77 | def draw(self, screen):
78 | pygame.surfarray.blit_array(screen, self.colors)
79 |
80 |
81 | def map_to_rainbow(in_array, out_array, low, high):
82 | if not SIMPLE_COLORING:
83 | # start at blue and rotate backwards
84 | # temp = (240 - numpy.round((in_array - low) * (360 / (high - low)))) % 360
85 | temp = in_array - low
86 | numpy.multiply(temp, 360 / (high - low), out=temp)
87 | numpy.round(temp, out=temp)
88 | numpy.subtract(240, temp, out=temp)
89 | numpy.remainder(temp, 360, out=temp)
90 |
91 | hues_to_rgb(temp, out_array)
92 | else:
93 | out_array[:] = numpy.round((in_array - low) / (high - low) * 63) * 0x020401
94 |
95 |
96 | def hues_to_rgb(h, out):
97 | X = numpy.zeros(h.shape, dtype=numpy.int32)
98 | numpy.multiply(255, 1 - numpy.abs((h / 60) % 2 - 1), out=X, casting='unsafe')
99 |
100 | r = numpy.zeros(h.shape, dtype=numpy.int32)
101 | g = numpy.zeros(h.shape, dtype=numpy.int32)
102 | b = numpy.zeros(h.shape, dtype=numpy.int32)
103 |
104 | h_lt_60 = h < 60
105 | h_bt_60_120 = (60 <= h) & (h < 120)
106 | h_bt_120_180 = (120 <= h) & (h < 180)
107 | h_bt_180_240 = (180 <= h) & (h < 240)
108 | h_bt_240_300 = (240 <= h) & (h < 300)
109 | h_gt_300 = 300 <= h
110 |
111 | r[h_lt_60] = 255
112 | r[h_bt_60_120] = X[h_bt_60_120]
113 | # r[h_bt_120_180] = 0 no-ops
114 | # r[h_bt_180_240] = 0
115 | r[h_bt_240_300] = X[h_bt_240_300]
116 | r[h_gt_300] = 255
117 |
118 | g[h_lt_60] = X[h_lt_60]
119 | g[h_bt_60_120] = 255
120 | g[h_bt_120_180] = 255
121 | g[h_bt_180_240] = X[h_bt_180_240]
122 | # g[h_bt_240_300] = 0
123 | # g[h_gt_300] = 0
124 |
125 | # b[h_lt_60] = 0
126 | # b[h_bt_60_120] = 0
127 | b[h_bt_120_180] = X[h_bt_120_180]
128 | b[h_bt_180_240] = 255
129 | b[h_bt_240_300] = 255
130 | b[h_gt_300] = X[h_gt_300]
131 |
132 | out[:] = r * 0x010000 + g * 0x000100 + b * 0x000001
133 |
134 |
135 | if __name__ == "__main__":
136 | pygame.init()
137 | screen = pygame.display.set_mode((W, H), flags=FLAGS)
138 | state = State()
139 | clock = pygame.time.Clock()
140 |
141 | running = True
142 | last_update_time = pygame.time.get_ticks()
143 |
144 | while running:
145 | current_time = pygame.time.get_ticks()
146 | for e in pygame.event.get():
147 | if e.type == pygame.QUIT or (e.type == pygame.KEYDOWN and e.key == pygame.K_ESCAPE):
148 | running = False
149 | elif e.type == pygame.KEYDOWN:
150 | if e.key == pygame.K_r:
151 | print("Resetting.")
152 | state = State()
153 | elif e.key == pygame.K_b:
154 | SHOW_BOX = not SHOW_BOX
155 | elif e.key == pygame.K_c:
156 | SIMPLE_COLORING = not SIMPLE_COLORING
157 | elif e.type == pygame.MOUSEBUTTONDOWN:
158 | r = state.rect
159 | xy = (r[0] + e.pos[0] / W * r[2], r[1] + e.pos[1] / H * r[3])
160 |
161 | zoom_change = None
162 | if e.button == 1:
163 | zoom_change = 2.0
164 | elif e.button == 3:
165 | zoom_change = 0.5
166 |
167 | if zoom_change is not None:
168 | new_rect = (xy[0] - r[2] / (2 * zoom_change),
169 | xy[1] - r[3] / (2 * zoom_change),
170 | r[2] / zoom_change,
171 | r[3] / zoom_change)
172 | print("Zooming to: {}".format(new_rect))
173 | state = State(view_rect=new_rect)
174 | state.update()
175 | state.draw(screen)
176 |
177 | mouse_xy = pygame.mouse.get_pos()
178 | if SHOW_BOX and 0 < mouse_xy[0] < W - 1 and 0 < mouse_xy[1] < H - 1:
179 | rect = pygame.Rect(mouse_xy[0] - W // 4, mouse_xy[1] - H // 4, W // 2, H // 2)
180 | pygame.draw.rect(screen, (255, 255, 255), rect, width=1)
181 |
182 | pygame.display.flip()
183 |
184 | if current_time // 250 > last_update_time // 250:
185 | pygame.display.set_caption("Fractal (FPS={:.1f}, SIZE={}, ITERS={}, VIEW={})".format(
186 | clock.get_fps(), (W, H), state.get_deepest_iteration(), state.get_pretty_rect()))
187 |
188 | last_update_time = current_time
189 | clock.tick(FPS)
190 |
--------------------------------------------------------------------------------
/gifreader.py:
--------------------------------------------------------------------------------
1 | import pygame
2 | import imageio.v3 as iio
3 | import typing
4 |
5 | import numpy
6 |
7 |
8 | def load_gif(filename) -> typing.Tuple[typing.List[pygame.Surface],
9 | typing.Dict[str, typing.Any]]:
10 | """Loads the image data and metadata from a gif file.
11 |
12 | Args:
13 | filename: The gif file to load.
14 |
15 | Returns:
16 | A tuple containing two items:
17 | output[0]: A list of Surfaces corresponding to the frames of the gif.
18 | output[1]: A dict containing metadata about the gif. The keys are strings and will always be present:
19 | 'file': (str) The file the gif was loaded from (matches the filename input).
20 | 'width': (int) The width of the image frames in pixels.
21 | 'height': (int) The height of the image frames in pixels.
22 | 'dims': (int, int) The dimensions of the image frames (width, height) in pixels.
23 | 'length': (int) The number of frames in the gif.
24 | 'duration': (int) The duration of each frame in the gif, in milliseconds.
25 | """
26 | # read gif metadata
27 | meta = iio.immeta(f"{filename}", extension=".gif", index=None)
28 |
29 | # ndarray with (num_frames, height, width, channel)
30 | gif = iio.imread(f"{filename}", extension=".gif", index=None)
31 | gif = numpy.transpose(gif, axes=(0, 2, 1, 3)) # flip x and y axes
32 |
33 | return [pygame.surfarray.make_surface(gif[i]) for i in range(gif.shape[0])], {
34 | 'file': filename,
35 | 'width': gif.shape[1],
36 | 'height': gif.shape[2],
37 | 'size': (gif.shape[1], gif.shape[2]),
38 | 'length': gif.shape[0],
39 | 'duration': int(meta['duration']) if 'duration' in meta else 0,
40 | 'loop': int(meta['loop']) if 'loop' in meta else 0
41 | }
42 |
43 |
44 | if __name__ == "__main__":
45 | FILE = "data/MOSHED-doctor.gif" # your gif goes here
46 | TEXT_SIZE = 32
47 |
48 | pygame.init()
49 | pygame.display.set_mode((640, 480), pygame.RESIZABLE)
50 | pygame.display.set_caption("gifreader.py")
51 |
52 | frames, metadata = load_gif(FILE)
53 | frm_duration = max(metadata['duration'], 1)
54 |
55 | font = pygame.font.Font(None, TEXT_SIZE)
56 |
57 | clock = pygame.time.Clock()
58 | elapsed_time = 0
59 |
60 | while True:
61 | for e in pygame.event.get():
62 | if e.type == pygame.QUIT:
63 | raise SystemExit()
64 |
65 | screen = pygame.display.get_surface()
66 | screen.fill((0, 0, 0))
67 | screen_size = screen.get_size()
68 |
69 | idx = (elapsed_time // frm_duration) % len(frames)
70 |
71 | offs = (screen_size[0] // 2 - frames[idx].get_width() // 2,
72 | screen_size[1] // 2 - frames[idx].get_height() // 2)
73 | screen.blit(frames[idx], offs)
74 |
75 | if TEXT_SIZE > 0: # render metadata
76 | x, y = 0, 0
77 | for key in metadata:
78 | screen.blit(font.render(f"{key}: {metadata[key]}", True, (255, 255, 255), (0, 0, 0, 0)), (x, y))
79 | y += font.get_height()
80 | screen.blit(font.render(f"frame: {idx}", True, (255, 255, 255), (0, 0, 0, 0)), (x, y))
81 | y += font.get_height()
82 |
83 | pygame.display.flip()
84 | elapsed_time += clock.tick(60)
85 |
--------------------------------------------------------------------------------
/life.py:
--------------------------------------------------------------------------------
1 | import pygame
2 | import numpy
3 |
4 | # Conway's Game of Life (using pygame + numpy)
5 | # by Ghast
6 |
7 | W, H = 1000, 600
8 | FPS = 60
9 | RATIO = 0.5 # Portion of cells that start out alive.
10 |
11 | FLAGS = 0 | pygame.SCALED
12 |
13 | class State:
14 | def __init__(self):
15 | self.grid = numpy.random.randint(0, 100, (W, H), numpy.int16)
16 | self.grid[self.grid <= RATIO * 100] = 1
17 | self.grid[self.grid > RATIO * 100] = 0
18 |
19 | self.neighbor_counts = numpy.zeros((W, H), numpy.int16)
20 |
21 | def step(self):
22 | self.neighbor_counts[:] = 0
23 | for dx in [-1, 0, 1]:
24 | for dy in [-1, 0, 1]:
25 | if (dx, dy) != (0, 0):
26 | shifted = numpy.roll(self.grid, (dx, dy), (0, 1))
27 | numpy.add(self.neighbor_counts, shifted, out=self.neighbor_counts)
28 |
29 | alive = self.grid > 0
30 | two = self.neighbor_counts == 2
31 | three = self.neighbor_counts == 3
32 |
33 | self.grid[:] = 0 # zero out the current grid
34 |
35 | # lots of unnecessary copying here, but it's so elegant
36 | self.grid[(alive & (two | three)) | ((~alive) & three)] = 1
37 |
38 | def draw(self, screen, color=0xFFFFFF):
39 | pygame.surfarray.blit_array(screen, self.grid * color)
40 |
41 | if __name__ == "__main__":
42 | pygame.init()
43 | screen = pygame.display.set_mode((W, H), flags=FLAGS)
44 | state = State()
45 | clock = pygame.time.Clock()
46 |
47 | running = True
48 | last_update_time = pygame.time.get_ticks()
49 |
50 | while running:
51 | current_time = pygame.time.get_ticks()
52 | dt = (current_time - last_update_time) / 1000
53 | for e in pygame.event.get():
54 | if e.type == pygame.QUIT or (e.type == pygame.KEYDOWN and e.key == pygame.K_ESCAPE):
55 | running = False
56 | elif e.type == pygame.KEYDOWN:
57 | if e.key == pygame.K_r:
58 | print("Reseting.")
59 | state = State()
60 | elif e.key == pygame.K_RIGHT:
61 | FPS += 10
62 | print("Increased target FPS to: {}".format(FPS))
63 | elif e.key == pygame.K_LEFT:
64 | FPS = max(10, FPS - 10)
65 | print("Decreased target FPS to: {}".format(FPS))
66 | state.step()
67 | state.draw(screen, color=0xAAFFAA)
68 |
69 | pygame.display.flip()
70 |
71 | if current_time // 1000 > last_update_time // 1000:
72 | pygame.display.set_caption("Game of Life (FPS={:.1f}, TARGET_FPS={}, SIZE={})".format(clock.get_fps(), FPS, (W, H)))
73 |
74 | last_update_time = current_time
75 | clock.tick(FPS)
76 |
--------------------------------------------------------------------------------
/lut.py:
--------------------------------------------------------------------------------
1 | import pygame
2 |
3 | try:
4 | import numpy # required for pygame.surfarray stuff
5 | _has_numpy = True
6 | except ImportError:
7 | _has_numpy = False
8 |
9 |
10 | _CACHE = {}
11 |
12 |
13 | def apply_lut(source: pygame.Surface, lut: pygame.Surface, idx: int, cache=True) -> pygame.Surface:
14 | """Generates a new copy of a surface with a color lookup table applied.
15 | source: Base image.
16 | lut: The "lookup table" surface. The first column of pixels should match the colors in the base image, and
17 | subsequent columns should contain different re-mappings for those colors.
18 | idx: Which column of the lookup table to use. If 0, colors will not be changed (since that's the "key column").
19 | cache: Whether to cache the result for quick lookup later. Note that this assumes the pixel values in the source
20 | and lut surfaces won't change between calls.
21 | """
22 | # this keying system assumes the source and lut's pixel values are static
23 | cache_key = (id(source), id(lut), idx) if cache else None
24 |
25 | if cache_key in _CACHE:
26 | return _CACHE[cache_key]
27 | else:
28 | if _has_numpy:
29 | res = _numpy_lut_iterative(source, lut, idx)
30 | else:
31 | res = _basic_pygame_lut(source, lut, idx)
32 |
33 | if cache_key is not None:
34 | _CACHE[cache_key] = res
35 | return res
36 |
37 |
38 | def _numpy_lut_iterative(source: pygame.Surface, lut: pygame.Surface, idx: int):
39 | # translucency isn't supported in lut
40 | if lut.get_flags() & pygame.SRCALPHA:
41 | lut = lut.convert()
42 |
43 | if source.get_flags() & pygame.SRCALPHA:
44 | # preserve source's alpha channel if it has one
45 | orig_alpha_array = pygame.surfarray.pixels_alpha(source)
46 | source = source.convert()
47 | else:
48 | orig_alpha_array = None
49 |
50 | res = source.copy()
51 | res_array = pygame.surfarray.pixels2d(res)
52 |
53 | source_array = pygame.surfarray.pixels2d(source)
54 | lut_array = pygame.surfarray.pixels2d(lut)
55 |
56 | # do the actual color swapping
57 | for y in range(lut.get_height()):
58 | res_array[source_array == lut_array[0, y]] = lut_array[idx, y]
59 |
60 | # restore alpha channel, if necessary
61 | if orig_alpha_array is not None:
62 | res = res.convert_alpha()
63 | res_alpha = pygame.surfarray.pixels_alpha(res)
64 | res_alpha[:] = orig_alpha_array
65 |
66 | return res
67 |
68 |
69 | def _basic_pygame_lut(source: pygame.Surface, lut: pygame.Surface, idx: int):
70 | res = source.copy()
71 |
72 | table = {}
73 | for y in range(lut.get_height()):
74 | table[tuple(lut.get_at((0, y)))] = lut.get_at((idx, y))
75 |
76 | for y in range(res.get_height()):
77 | for x in range(res.get_width()):
78 | c = tuple(res.get_at((x, y)))
79 | if c in table:
80 | res.set_at((x, y), table[c])
81 |
82 | return res
83 |
84 |
85 | if __name__ == "__main__":
86 | pygame.init()
87 |
88 | FPS = 60
89 |
90 | screen = pygame.display.set_mode((640, 480), flags=pygame.RESIZABLE)
91 | clock = pygame.time.Clock()
92 |
93 | BASE_IMG: pygame.Surface = None
94 | LUT_IMG: pygame.Surface = None
95 | LUT_IDX = 1
96 | NUM_LUTS = 2
97 |
98 | def load_sample_images():
99 | global BASE_IMG, LUT_IMG, LUT_IDX, NUM_LUTS
100 | BASE_IMG = pygame.image.load("data/frog_src.png").convert()
101 | BASE_IMG.set_colorkey(BASE_IMG.get_at((0, 0)))
102 |
103 | LUT_IMG = pygame.image.load("data/frog_lut.png").convert()
104 | NUM_LUTS = LUT_IMG.get_width()
105 | LUT_IDX = min(LUT_IDX, NUM_LUTS - 1)
106 |
107 | load_sample_images()
108 |
109 | running = True
110 | last_update_time = pygame.time.get_ticks()
111 |
112 | def _draw_in_rect(surf, img, rect: pygame.Rect):
113 | if img.get_height() / img.get_width() >= rect[3] / rect[2]:
114 | xformed = pygame.transform.scale(img, (img.get_width() * rect[3] / img.get_height(), rect[3]))
115 | else:
116 | xformed = pygame.transform.scale(img, (rect[2], img.get_height() * rect[2] / img.get_width()))
117 | blit_rect = xformed.get_rect(center=rect.center)
118 | surf.blit(xformed, blit_rect)
119 | return blit_rect
120 |
121 | while running:
122 | current_time = pygame.time.get_ticks()
123 | dt = (current_time - last_update_time) / 1000
124 | pygame.time.get_ticks()
125 | for e in pygame.event.get():
126 | if e.type == pygame.QUIT or (e.type == pygame.KEYDOWN and e.key == pygame.K_ESCAPE):
127 | running = False
128 | elif e.type == pygame.KEYDOWN:
129 | if e.key == pygame.K_RIGHT:
130 | LUT_IDX += 1
131 | elif e.key == pygame.K_LEFT:
132 | LUT_IDX -= 1
133 | elif e.key == pygame.K_r:
134 | print("INFO: reloading images...")
135 | load_sample_images()
136 |
137 | LUT_IDX %= NUM_LUTS
138 | RESULT_IMG = apply_lut(BASE_IMG, LUT_IMG, LUT_IDX)
139 |
140 | screen = pygame.display.get_surface()
141 | W, H = screen.get_size()
142 | screen.fill((66, 66, 66))
143 |
144 | to_draw = [BASE_IMG, LUT_IMG, RESULT_IMG]
145 |
146 | for i, img in enumerate(to_draw):
147 | rect = pygame.Rect(i * W / 3, 0, int(W / 3), H)
148 | blit_rect = _draw_in_rect(screen, img, rect)
149 | if img == LUT_IMG:
150 | selection_rect = pygame.Rect(blit_rect.x + LUT_IDX * blit_rect.width / NUM_LUTS,
151 | blit_rect.y,
152 | blit_rect.width / NUM_LUTS,
153 | blit_rect.height)
154 | pygame.draw.rect(screen, (255, 0, 0), selection_rect, width=2)
155 |
156 | pygame.display.flip()
157 |
158 | if current_time // 1000 > last_update_time // 1000:
159 | pygame.display.set_caption(
160 | "LUT Demo (FPS={:.1f}, TARGET_FPS={})".format(clock.get_fps(), FPS))
161 |
162 | last_update_time = current_time
163 | clock.tick(FPS)
164 |
--------------------------------------------------------------------------------
/pendulum.py:
--------------------------------------------------------------------------------
1 | import numpy
2 | import random
3 | import math
4 | import argparse
5 |
6 | import os
7 | os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
8 | import pygame
9 | import pygame.gfxdraw
10 |
11 | # Double Pendulum with Pygame + pyOpenGL
12 | # by Ghast ~ https://github.com/davidpendergast
13 |
14 | USE_GL = True # if False, will use pure pygame
15 |
16 | if USE_GL:
17 | import OpenGL as OpenGL
18 | import OpenGL.GL as GL
19 | import OpenGL.GLU as GLU
20 |
21 | # Params (can be set via the command line)
22 | N = 1000 # Number of pendulums
23 | L = 20 # Length of pendulum arms
24 | M = 5 # Mass of pendulum arms
25 | G = 2 * 9.8 # Gravity
26 | FPS = 60
27 | ZOOM = 5
28 | SPREAD = (6.283 / 360) / 4
29 |
30 | ML2 = M * L * L
31 |
32 | # OpenGL-Only Params
33 | COLOR_CHANNELS = 4 # must be 3 or 4
34 | RAINBOW = True
35 | OPACITY = 0.1
36 |
37 |
38 | def get_initial_conds():
39 | theta1 = 3.1415 * (random.random() + 0.5)
40 | theta2 = random.random() * 6.283
41 | return theta1, theta2, SPREAD
42 |
43 |
44 | class State:
45 |
46 | def __init__(self):
47 | theta1, theta2, spread = get_initial_conds()
48 | self.theta1 = numpy.linspace(theta1, theta1 + spread, N)
49 | self.theta2 = numpy.linspace(theta2, theta2 + spread, N)
50 | self.p1 = numpy.zeros(N)
51 | self.p2 = numpy.zeros(N)
52 |
53 | self.colors = [hsv_to_rgb(360 * i / N, 0.75, 1) for i in range(N)]
54 | self.colors.reverse()
55 |
56 | # arrays for temp storage (to avoid reallocating arrays constantly)
57 | self.sub = numpy.zeros(N)
58 | self.cos = numpy.zeros(N)
59 | self.cos2 = numpy.zeros(N)
60 | self.sin = numpy.zeros(N)
61 | self.temp1 = numpy.zeros(N)
62 | self.temp2 = numpy.zeros(N)
63 | self.temp3 = numpy.zeros(N)
64 | self.temp4 = numpy.zeros(N)
65 | self.denom = numpy.zeros(N)
66 | self.num = numpy.zeros(N)
67 |
68 | self.dtheta1 = numpy.zeros(N)
69 | self.dtheta2 = numpy.zeros(N)
70 | self.dp1 = numpy.zeros(N)
71 | self.dp2 = numpy.zeros(N)
72 | self.x1 = numpy.zeros(N, dtype=numpy.int16)
73 | self.y1 = numpy.zeros(N, dtype=numpy.int16)
74 | self.x2 = numpy.zeros(N, dtype=numpy.int16)
75 | self.y2 = numpy.zeros(N, dtype=numpy.int16)
76 |
77 | if USE_GL:
78 | self.vertex_data = numpy.zeros((N * 2 + 1) * 2, dtype=float)
79 | self.color_data = numpy.ones((N * 2 + 1) * COLOR_CHANNELS, dtype=float)
80 | if RAINBOW:
81 | for i in range(0, N * 2):
82 | c = self.colors[i // 2]
83 | self.color_data[i * COLOR_CHANNELS + 0] = c[0] / 256
84 | self.color_data[i * COLOR_CHANNELS + 1] = c[1] / 256
85 | self.color_data[i * COLOR_CHANNELS + 2] = c[2] / 256
86 | if COLOR_CHANNELS > 3:
87 | self.color_data[3::COLOR_CHANNELS] = OPACITY
88 |
89 | self.index_data = numpy.arange(0, N * 4, dtype=int)
90 | self.index_data[0::4] = N * 2 # center point is stored as the Nth vertex
91 | self.index_data[1::4] = numpy.arange(0, N * 2, 2, dtype=int)
92 | self.index_data[2::4] = numpy.arange(0, N * 2, 2, dtype=int)
93 | self.index_data[3::4] = numpy.arange(1, N * 2, 2, dtype=int)
94 |
95 | def euler_update(self, dt):
96 | numpy.subtract(self.theta1, self.theta2, out=self.sub)
97 | numpy.cos(self.sub, out=self.cos)
98 | numpy.square(self.cos, out=self.cos2)
99 | numpy.sin(self.sub, out=self.sin)
100 |
101 | self.calc_dtheta1(self.p1, self.p2, self.cos, self.cos2)
102 | self.calc_dtheta2(self.p1, self.p2, self.cos, self.cos2)
103 | self.calc_dp1(self.theta1, self.dtheta1, self.dtheta2, self.sin)
104 | self.calc_dp2(self.theta2, self.dtheta1, self.dtheta2, self.sin)
105 |
106 | self.theta1 = self.theta1 + dt * self.dtheta1
107 | self.theta2 = self.theta2 + dt * self.dtheta2
108 | self.p1 = self.p1 + dt * self.dp1
109 | self.p2 = self.p2 + dt * self.dp2
110 |
111 | def calc_dtheta1(self, p1, p2, cos, cos2):
112 | # self.dtheta1 = (6 / ML2) * (2 * p1 - 3 * cos * p2) / (16 - 9 * cos2)
113 | numpy.multiply(2, p1, out=self.temp1)
114 | numpy.multiply(3, cos, out=self.temp2)
115 | numpy.multiply(self.temp2, p2, out=self.temp2)
116 | numpy.subtract(self.temp1, self.temp2, out=self.num)
117 | numpy.multiply((6 / ML2), self.num, out=self.num)
118 | numpy.multiply(9, cos2, out=self.temp3)
119 | numpy.subtract(16, self.temp3, out=self.denom)
120 | numpy.divide(self.num, self.denom, out=self.dtheta1)
121 |
122 | def calc_dtheta2(self, p1, p2, cos, cos2):
123 | # self.dtheta2 = (6 / ML2) * (8 * p2 - 3 * cos * p1) / (16 - 9 * cos2)
124 | numpy.multiply(8, p2, out=self.temp1)
125 | numpy.multiply(3, cos, out=self.temp2)
126 | numpy.multiply(self.temp2, p1, out=self.temp2)
127 | numpy.subtract(self.temp1, self.temp2, out=self.num)
128 | numpy.multiply((6 / ML2), self.num, out=self.num)
129 | numpy.multiply(9, cos2, out=self.temp3)
130 | numpy.subtract(16, self.temp3, out=self.denom)
131 | numpy.divide(self.num, self.denom, out=self.dtheta2)
132 |
133 | def calc_dp1(self, theta1, dtheta1, dtheta2, sin):
134 | # self.dp1 = (-ML2 / 2) * (dtheta1 * dtheta2 * sin + 3 * G / L * numpy.sin(theta1))
135 | numpy.multiply(dtheta1, dtheta2, out=self.temp1)
136 | numpy.multiply(self.temp1, sin, out=self.temp1)
137 | numpy.sin(theta1, out=self.temp2)
138 | numpy.multiply(3 * G / L, self.temp2, out=self.temp2)
139 | numpy.add(self.temp1, self.temp2, out=self.temp3)
140 | numpy.multiply(-ML2 / 2, self.temp3, out=self.dp1)
141 |
142 | def calc_dp2(self, theta2, dtheta1, dtheta2, sin):
143 | # self.dp2 = (-ML2 / 2) * (-dtheta1 * dtheta2 * sin + G / L * numpy.sin(theta2))
144 | numpy.multiply(dtheta1, dtheta2, out=self.temp1)
145 | numpy.multiply(self.temp1, sin, out=self.temp1)
146 | numpy.sin(theta2, out=self.temp2)
147 | numpy.multiply(G / L, self.temp2, out=self.temp2)
148 | numpy.subtract(self.temp2, self.temp1, out=self.temp3)
149 | numpy.multiply(-ML2 / 2, self.temp3, out=self.dp2)
150 |
151 | def render_all(self, screen):
152 | x0 = screen.get_width() // 2
153 | y0 = screen.get_height() // 2
154 |
155 | # self.x1 = x0 + ZOOM * L * numpy.cos(self.theta1 + 3.1429 / 2)
156 | numpy.add(self.theta1, 3.1429 / 2, out=self.temp1)
157 | numpy.cos(self.temp1, out=self.temp1)
158 | numpy.multiply(ZOOM * L, self.temp1, out=self.temp1)
159 | numpy.add(x0, self.temp1, out=self.x1, casting='unsafe')
160 |
161 | # self.y1 = y0 + ZOOM * L * numpy.sin(self.theta1 + 3.1429 / 2)
162 | numpy.add(self.theta1, 3.1429 / 2, out=self.temp2)
163 | numpy.sin(self.temp2, out=self.temp2)
164 | numpy.multiply(ZOOM * L, self.temp2, out=self.temp2)
165 | numpy.add(y0, self.temp2, out=self.y1, casting='unsafe')
166 |
167 | # self.x2 = self.x1 + ZOOM * L * numpy.cos(self.theta2 + 3.1429 / 2)
168 | numpy.add(self.theta2, 3.1429 / 2, out=self.temp3)
169 | numpy.cos(self.temp3, out=self.temp3)
170 | numpy.multiply(ZOOM * L, self.temp3, out=self.temp3)
171 | numpy.add(self.x1, self.temp3, out=self.x2, casting='unsafe')
172 |
173 | # self.y2 = self.y1 + ZOOM * L * numpy.sin(self.theta2 + 3.1429 / 2)
174 | numpy.add(self.theta2, 3.1429 / 2, out=self.temp4)
175 | numpy.sin(self.temp4, out=self.temp4)
176 | numpy.multiply(ZOOM * L, self.temp4, out=self.temp4)
177 | numpy.add(self.y1, self.temp4, out=self.y2, casting='unsafe')
178 |
179 | if USE_GL:
180 | self.vertex_data[0:N*4:4] = self.x1
181 | self.vertex_data[1:N*4:4] = self.y1
182 | self.vertex_data[2:N*4:4] = self.x2
183 | self.vertex_data[3:N*4:4] = self.y2
184 | self.vertex_data[N*4] = x0
185 | self.vertex_data[N*4 + 1] = y0
186 |
187 | GL.glClear(GL.GL_COLOR_BUFFER_BIT)
188 |
189 | GL.glEnableClientState(GL.GL_VERTEX_ARRAY)
190 | GL.glEnableClientState(GL.GL_COLOR_ARRAY)
191 |
192 | GL.glVertexPointer(2, GL.GL_FLOAT, 0, self.vertex_data)
193 | GL.glColorPointer(COLOR_CHANNELS, GL.GL_FLOAT, 0, self.color_data)
194 | GL.glDrawElements(GL.GL_LINES, len(self.index_data), GL.GL_UNSIGNED_INT, self.index_data);
195 |
196 | GL.glDisableClientState(GL.GL_VERTEX_ARRAY)
197 | GL.glDisableClientState(GL.GL_COLOR_ARRAY)
198 |
199 | else:
200 | screen.fill((0, 0, 0))
201 | # (-_-) don't think there's a good way to avoid this loop without gl...
202 | for i in range(0, N):
203 | pygame.gfxdraw.line(screen, x0, y0, self.x1[i], self.y1[i], self.colors[i])
204 | pygame.gfxdraw.line(screen, self.x1[i], self.y1[i], self.x2[i], self.y2[i], self.colors[i])
205 |
206 |
207 | import cProfile
208 | import pstats
209 |
210 |
211 | class Profiler:
212 |
213 | def __init__(self):
214 | self.is_running = False
215 | self.pr = cProfile.Profile(builtins=False)
216 |
217 | def toggle(self):
218 | self.is_running = not self.is_running
219 |
220 | if not self.is_running:
221 | self.pr.disable()
222 |
223 | ps = pstats.Stats(self.pr)
224 | ps.strip_dirs()
225 | ps.sort_stats('cumulative')
226 | ps.print_stats(35)
227 |
228 | else:
229 | print("Started profiling...")
230 | self.pr.clear()
231 | self.pr.enable()
232 |
233 |
234 | def initialize_display(size):
235 | pygame.init()
236 |
237 | if USE_GL:
238 | display = size
239 | flags = pygame.DOUBLEBUF | pygame.OPENGL
240 | screen = pygame.display.set_mode(display, flags)
241 | GL.glClearColor(0, 0, 0, 1)
242 | GL.glViewport(0, 0, display[0], display[1])
243 | GL.glOrtho(0.0, display[0], display[1], 0.0, 0.0, 1.0);
244 | if COLOR_CHANNELS > 3:
245 | GL.glEnable(GL.GL_BLEND)
246 | GL.glBlendFunc(GL.GL_SRC_ALPHA, GL.GL_ONE_MINUS_SRC_ALPHA);
247 |
248 | return screen
249 | else:
250 | return pygame.display.set_mode(size, 0, 8)
251 |
252 |
253 | def start(size):
254 | screen = initialize_display(size)
255 | clock = pygame.time.Clock()
256 |
257 | state = State()
258 | profiler = Profiler()
259 |
260 | running = True
261 | last_update_time = pygame.time.get_ticks()
262 |
263 | while running:
264 | current_time = pygame.time.get_ticks()
265 | dt = (current_time - last_update_time) / 1000
266 | state.euler_update(dt)
267 |
268 | state.render_all(screen)
269 | pygame.display.flip()
270 |
271 | for event in pygame.event.get():
272 | if event.type == pygame.QUIT or (event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE):
273 | running = False
274 | elif event.type == pygame.KEYDOWN:
275 | if event.key == pygame.K_r:
276 | print("Reseting...")
277 | state = State()
278 | elif event.key == pygame.K_p:
279 | profiler.toggle()
280 |
281 | if current_time // 1000 > last_update_time // 1000:
282 | pygame.display.set_caption("Double Pendulum (FPS={:.1f}, N={})".format(clock.get_fps(), N))
283 |
284 | last_update_time = current_time
285 | clock.tick(FPS)
286 |
287 |
288 | def hsv_to_rgb(h, s, v):
289 | """
290 | :param h: 0 <= h < 360
291 | :param s: 0 <= s <= 1
292 | :param v: 0 <= v <= 1
293 | :return: (r, g, b) as floats
294 | """
295 | C = v * s
296 | X = C * (1 - abs((h / 60) % 2 - 1))
297 | m = v - C
298 |
299 | if h < 60:
300 | rgb_prime = (C, X, 0)
301 | elif h < 120:
302 | rgb_prime = (X, C, 0)
303 | elif h < 180:
304 | rgb_prime = (0, C, X)
305 | elif h < 240:
306 | rgb_prime = (0, X, C)
307 | elif h < 300:
308 | rgb_prime = (X, 0, C)
309 | else:
310 | rgb_prime = (C, 0, X)
311 |
312 | return (int(256 * rgb_prime[0] + m),
313 | int(256 * rgb_prime[1] + m),
314 | int(256 * rgb_prime[2] + m))
315 |
316 |
317 | if __name__ == "__main__":
318 | parser = argparse.ArgumentParser(description='A double pendulum simulation made with pygame, PyOpenGL, and numpy.')
319 | parser.add_argument('-n', type=int, metavar="int", default=N, help=f'the number of pendulums (default {N})')
320 | parser.add_argument('--opacity', type=float, metavar="float", default=OPACITY, help=f'the opacity of the pendulums (default {OPACITY})')
321 | parser.add_argument('--length', type=int, metavar="int", default=L, help=f'the length of the pendulum arms (default {L})')
322 | parser.add_argument('--mass', type=int, metavar="float", default=M, help=f'the mass of the pendulum arms (default {M})')
323 | parser.add_argument('--fps', type=int, metavar="int", default=FPS, help=f'the target FPS for the simulation (default {FPS})')
324 | parser.add_argument('--zoom', type=int, metavar="int", default=ZOOM, help=f'the target FPS for the simulation (default {ZOOM})')
325 | parser.add_argument('--size', type=int, metavar="int", default=[800, 600], nargs=2, help='the window size for the simulation (default 800 600)')
326 | parser.add_argument('--spread', type=float, metavar="float", default=SPREAD, help=f'the initial spread, in radians (default {SPREAD:.4f})')
327 |
328 | args = parser.parse_args()
329 |
330 | N = args.n
331 | OPACITY = args.opacity
332 | L = args.length
333 | M = args.mass
334 | FPS = args.fps
335 | ZOOM = args.zoom
336 | SPREAD = args.spread
337 |
338 | start(args.size)
339 |
--------------------------------------------------------------------------------
/rainbowize.py:
--------------------------------------------------------------------------------
1 |
2 | import math
3 | import numpy
4 |
5 | import pygame
6 | import pygame._sdl2 as sdl2
7 |
8 | import os
9 |
10 |
11 | def apply_rainbow(surface: pygame.Surface, offset=0., strength=0.666, bands=2.) -> pygame.Surface:
12 | """Adds a rainbow effect to an image.
13 |
14 | Note that this returns a new surface and does not modify the original.
15 |
16 | Args:
17 | surface: The original image.
18 | offset: A value from 0 to 1 that applies a color shift to the rainbow. Changing this parameter
19 | over time will create an animated looping effect. Values outside the interval [0, 1) will
20 | act as though they're modulo 1.
21 | strength: A value from 0 to 1 that determines how strongly the rainbow's color should appear.
22 | bands: A value greater than 0 that determines how many color bands should be rendered (in other words,
23 | how "thin" the rainbow should appear).
24 | """
25 | x = numpy.linspace(0, 1, surface.get_width())
26 | y = numpy.linspace(0, 1, surface.get_height())
27 | gradient = numpy.outer(x, y) * bands
28 |
29 | red_mult = numpy.sin(math.pi * 2 * (gradient + offset)) * 0.5 + 0.5
30 | green_mult = numpy.sin(math.pi * 2 * (gradient + offset + 0.25)) * 0.5 + 0.5
31 | blue_mult = numpy.sin(math.pi * 2 * (gradient + offset + 0.5)) * 0.5 + 0.5
32 |
33 | res = surface.copy()
34 |
35 | red_pixels = pygame.surfarray.pixels_red(res)
36 | red_pixels[:] = (red_pixels * (1 - strength) + red_pixels * red_mult * strength).astype(dtype='uint16')
37 |
38 | green_pixels = pygame.surfarray.pixels_green(res)
39 | green_pixels[:] = (green_pixels * (1 - strength) + green_pixels * green_mult * strength).astype(dtype='uint16')
40 |
41 | blue_pixels = pygame.surfarray.pixels_blue(res)
42 | blue_pixels[:] = (blue_pixels * (1 - strength) + blue_pixels * blue_mult * strength).astype(dtype='uint16')
43 |
44 | return res
45 |
46 |
47 | def make_fancy_scaled_display(
48 | size,
49 | scale_factor=0.,
50 | extra_flags=0,
51 | outer_fill_color=None,
52 | smooth: bool = None) -> pygame.Surface:
53 | """Creates a SCALED pygame display with some additional customization options.
54 |
55 | Args:
56 | size: The base resolution of the display surface.
57 | extra_flags: Extra display flags (aside from SCALED) to give the display, e.g. pygame.RESIZABLE.
58 | scale_factor: The initial scaling factor for the window.
59 | For example, if the display's base size is 64x32 and this arg is 5, the window will be 320x160
60 | in the physical display. If this arg is 0 or less, the window will use the default SCALED behavior
61 | of filling as much space as the computer's display will allow.
62 | Non-integer values greater than 1 can be used here too. Positive values less than 1 will act like 1.
63 | outer_fill_color: When the display surface can't cleanly fill the physical window with an integer scale
64 | factor, a solid color is used to fill the empty space. If provided, this param sets that color
65 | (otherwise it's black by default).
66 | smooth: Whether to use smooth interpolation while scaling.
67 | If True: The environment variable PYGAME_FORCE_SCALE will be set to 'photo', which according to
68 | the pygame docs, "makes the scaling use the slowest, but highest quality anisotropic scaling
69 | algorithm, if it is available." This gives a smoother, blurrier look.
70 | If False: PYGAME_FORCE_SCALE will be set to 'default', which uses nearest-neighbor interpolation.
71 | If None: PYGAME_FORCE_SCALE is left unchanged, resulting in nearest-neighbor interpolation (unless
72 | the variable has been set beforehand). This is the default behavior.
73 | Returns:
74 | The display surface.
75 | """
76 |
77 | if smooth is not None:
78 | # must be set before display.set_mode is called.
79 | os.environ['PYGAME_FORCE_SCALE'] = 'photo' if smooth else 'default'
80 |
81 | # create the display in "hidden" mode, because it isn't properly sized yet
82 | res = pygame.display.set_mode(size, pygame.SCALED | extra_flags | pygame.HIDDEN)
83 |
84 | window = sdl2.Window.from_display_module()
85 |
86 | # due to a bug, we *cannot* let this Window object get GC'd
87 | # https://github.com/pygame-community/pygame-ce/issues/1889
88 | globals()["sdl2_Window_ref"] = window # so store it somewhere safe...
89 |
90 | # resize the window to achieve the desired scale factor
91 | if scale_factor > 0:
92 | initial_scale_factor = max(scale_factor, 1) # scale must be >= 1
93 | window.size = (int(size[0] * initial_scale_factor),
94 | int(size[1] * initial_scale_factor))
95 | window.position = sdl2.WINDOWPOS_CENTERED # recenter it too
96 |
97 | # set the out-of-bounds color
98 | if outer_fill_color is not None:
99 | renderer = sdl2.Renderer.from_window(window)
100 | renderer.draw_color = pygame.Color(outer_fill_color)
101 |
102 | # show the window (unless the HIDDEN flag was passed in)
103 | if not (pygame.HIDDEN & extra_flags):
104 | window.show()
105 |
106 | return res
107 |
108 |
109 | if __name__ == "__main__":
110 | pygame.init()
111 |
112 | bg_color = pygame.Color(92, 64, 92)
113 | outer_bg_color = bg_color.lerp("black", 0.25)
114 |
115 | screen = make_fancy_scaled_display(
116 | (256, 128),
117 | extra_flags=pygame.RESIZABLE,
118 | scale_factor=3,
119 | outer_fill_color=outer_bg_color,
120 | smooth=False # looks bad for pixel art
121 | )
122 |
123 | pygame.display.set_caption("rainbowize.py")
124 |
125 | frog_img = pygame.image.load("data/frog.png").convert_alpha()
126 | frog_img = pygame.transform.scale(frog_img, (frog_img.get_width() * 3,
127 | frog_img.get_height() * 3))
128 |
129 | clock = pygame.time.Clock()
130 |
131 | while True:
132 | for e in pygame.event.get():
133 | if e.type == pygame.QUIT:
134 | raise SystemExit()
135 | if e.type == pygame.KEYDOWN:
136 | if e.key == pygame.K_ESCAPE:
137 | raise SystemExit()
138 |
139 | screen.fill(bg_color)
140 |
141 | # loop the animation 1.5 times per second
142 | elapsed_time_ms = pygame.time.get_ticks()
143 | animation_period_ms = 666
144 |
145 | rainbow_frog_img = apply_rainbow(
146 | frog_img,
147 | offset=elapsed_time_ms / animation_period_ms,
148 | bands=1.2
149 | )
150 |
151 | scr_w, scr_h = screen.get_size()
152 | screen.blit(frog_img, (scr_w / 4 - frog_img.get_width() / 2, scr_h / 2 - frog_img.get_height() / 2))
153 | screen.blit(rainbow_frog_img, (3 * scr_w / 4 - rainbow_frog_img.get_width() / 2,
154 | scr_h / 2 - rainbow_frog_img.get_height() / 2))
155 |
156 | pygame.display.flip()
157 | clock.tick(60)
158 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | PyOpenGl>=3.1.0
2 | pygame>=2.0.1
3 | numpy>=1.16.4
4 | opencv-python>=4.5.2.54
5 | imageio>=2.22.4
6 |
--------------------------------------------------------------------------------
/screenshots/benchmark_results.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/benchmark_results.png
--------------------------------------------------------------------------------
/screenshots/colorswirl.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/colorswirl.gif
--------------------------------------------------------------------------------
/screenshots/fractal_zoom.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/fractal_zoom.gif
--------------------------------------------------------------------------------
/screenshots/gifreader_demo.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/gifreader_demo.gif
--------------------------------------------------------------------------------
/screenshots/life_demo.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/life_demo.gif
--------------------------------------------------------------------------------
/screenshots/lut.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/lut.gif
--------------------------------------------------------------------------------
/screenshots/n=1000.PNG:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/n=1000.PNG
--------------------------------------------------------------------------------
/screenshots/n=10000.PNG:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/n=10000.PNG
--------------------------------------------------------------------------------
/screenshots/n=100000.PNG:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/n=100000.PNG
--------------------------------------------------------------------------------
/screenshots/n=1000animated.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/n=1000animated.gif
--------------------------------------------------------------------------------
/screenshots/rainbowize.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/rainbowize.gif
--------------------------------------------------------------------------------
/screenshots/stega_demo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/stega_demo.png
--------------------------------------------------------------------------------
/screenshots/video_demo.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/video_demo.gif
--------------------------------------------------------------------------------
/screenshots/warp.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/warp.gif
--------------------------------------------------------------------------------
/screenshots/worms.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/davidpendergast/pygame-utils/86828475f68a3a37c7d1a24069a3a9af454d6d90/screenshots/worms.gif
--------------------------------------------------------------------------------
/steganography.py:
--------------------------------------------------------------------------------
1 | import math
2 | import typing
3 |
4 | import pygame
5 | import numpy
6 |
7 |
8 | def write_text_to_surface(text_data: str, input_surface: pygame.Surface,
9 | bit_depth_range: typing.Union[typing.Tuple[int, int], int] = (1, 4),
10 | end_str: str = chr(0),
11 | resize_mode: str = 'smooth') -> pygame.Surface:
12 | """Writes ascii data into the pixel values of a surface (RGB only).
13 |
14 | This function returns a new copy of the image, leaving the original unmodified. Lower-ordered bits of each color
15 | channel are consumed first, to preserve the image as best as possible. A lower bit depth will give better image
16 | quality at the cost of worse data compression, and vice-versa. Noise is added after the end of the data section
17 | to avoid creating an obvious boundary.
18 |
19 | Args:
20 | text_data: The ascii data to write.
21 | input_surface: The surface to write the data into.
22 | bit_depth_range: The range of bits to use for data storage, per byte of image data. An optimal value
23 | will be selected from this range based on the image's size and the amount of ascii data. Bounds must be
24 | between 1 and 8. If a single int is provided, that bit depth will be used.
25 | end_str: Indicator for where the data ends. This can be any string or character that doesn't appear
26 | in the data string. By default, it's the NULL character (0x00000000).
27 | resize_mode: This parameter controls how the function behaves when the image isn't large enough
28 | to fit the data. Accepted values: None, "smooth", or "integer". If None, the image will not be resized,
29 | and the function will throw an error if the data doesn't fit. If "smooth", the image will be scaled up
30 | smoothly in each axis. If "integer", the image will be scaled up by an integer multiplier in each axis.
31 |
32 | Returns:
33 | A copy of the input surface with the text data encoded into its pixel values.
34 |
35 | Raises:
36 | ValueError: If any parameters are invalid, or resize_mode is None and the data is too large to be stored.
37 | """
38 | if isinstance(bit_depth_range, int):
39 | bit_depth_range = (bit_depth_range, bit_depth_range)
40 | if not (1 <= bit_depth_range[0] <= bit_depth_range[1] <= 8):
41 | raise ValueError(f"Illegal bit_depth_range: {bit_depth_range}")
42 |
43 | if end_str in text_data:
44 | raise ValueError(f"text_data cannot contain end_str (found at index: {text_data.index(end_str)})")
45 | text_data += end_str
46 |
47 | bytes_in_img = input_surface.get_width() * input_surface.get_height() * 3
48 | if bytes_in_img == 0:
49 | raise ValueError("Cannot write text to empty surface.")
50 |
51 | # find an optimal bit depth within the specified bounds.
52 | header_data_size = 3
53 | optimal_bit_depth = math.ceil(len(text_data) * 8 / (bytes_in_img - header_data_size))
54 | bit_depth = max(bit_depth_range[0], min(optimal_bit_depth, bit_depth_range[1]))
55 |
56 | data_array = _str_to_flat_array(text_data, bits_per_index=bit_depth)
57 | img_bytes_needed = header_data_size + data_array.size
58 | img_bytes_in_input = input_surface.get_width() * input_surface.get_height() * 3
59 |
60 | # resize the input surface (if necessary) so it's large enough to hold the data.
61 | if img_bytes_in_input <= img_bytes_needed:
62 | if resize_mode == 'integer':
63 | mult = math.ceil(math.sqrt(img_bytes_needed / img_bytes_in_input))
64 | new_dims = (input_surface.get_width() * mult,
65 | input_surface.get_height() * mult)
66 | elif resize_mode == 'smooth':
67 | mult = math.sqrt(img_bytes_needed / img_bytes_in_input)
68 | new_dims = (math.ceil(input_surface.get_width() * mult),
69 | math.ceil(input_surface.get_height() * mult))
70 | else:
71 | raise ValueError(f"The surface is too small to contain {len(text_data)} bytes of text "
72 | f"with a bit_depth of {bit_depth}.")
73 | output_surface = pygame.transform.scale(input_surface, new_dims)
74 | else:
75 | output_surface = input_surface.copy()
76 |
77 | img_bytes_in_output = output_surface.get_width() * output_surface.get_height() * 3
78 | if img_bytes_needed < img_bytes_in_output:
79 | end_idx = header_data_size + data_array.size
80 | data_array = numpy.pad(data_array, (header_data_size, img_bytes_in_output - (data_array.size + header_data_size)),
81 | 'constant', constant_values=(0, 0))
82 | # fill the rest of the image with noise, to avoid creating a visible boundary.
83 | data_array[end_idx:] = numpy.random.randint(2 ** bit_depth, size=data_array.size - end_idx)
84 |
85 | # write 1 bit of 'header data' into the first 3 bytes of the image, to indicate the bit depth of the data section.
86 | first_px_rgb = list(output_surface.get_at((0, 0)))
87 | first_px_rgb[0] = _set_bit(first_px_rgb[0], 0, (bit_depth // 1) % 2)
88 | first_px_rgb[1] = _set_bit(first_px_rgb[1], 0, (bit_depth // 2) % 2)
89 | first_px_rgb[2] = _set_bit(first_px_rgb[2], 0, (bit_depth // 4) % 2)
90 |
91 | colors = [
92 | pygame.surfarray.pixels_red(output_surface),
93 | pygame.surfarray.pixels_green(output_surface),
94 | pygame.surfarray.pixels_blue(output_surface)
95 | ]
96 |
97 | # finally, write the actual data.
98 | for c in range(3):
99 | colors[c] &= 255 - ((1 << bit_depth) - 1) # e.g. 0x11110000, where # of zeros = bit_depth
100 | colors[c] |= data_array[c::3].reshape(colors[c].shape)
101 |
102 | # write header data (first 3 bytes = RGB channels of the 1st pixel in the image).
103 | output_surface.set_at((0, 0), first_px_rgb)
104 |
105 | return output_surface
106 |
107 |
108 | def read_text_from_surface(surface: pygame.Surface, end_str=chr(0)) -> str:
109 | """Extracts the ascii data that was written into a surface by write_text_to_surface(...).
110 | surface: The surface.
111 | end_str: Indicator for where the data ends. Must match the string that was used when writing the data.
112 | """
113 | # first, read the header data to find the bit_depth
114 | first_px_rgb = surface.get_at((0, 0))
115 | bit_depth = first_px_rgb[0] % 2 + (first_px_rgb[1] % 2) * 2 + (first_px_rgb[2] % 2) * 4
116 | if not (1 <= bit_depth <= 8):
117 | raise ValueError(f"Illegal bit_depth: {bit_depth}")
118 | header_data_size = 3
119 |
120 | colors = [
121 | pygame.surfarray.pixels_red(surface),
122 | pygame.surfarray.pixels_green(surface),
123 | pygame.surfarray.pixels_blue(surface)
124 | ]
125 |
126 | raw_data = numpy.array([0] * (surface.get_width() * surface.get_height() * 3), dtype="uint8")
127 | mask = (1 << bit_depth) - 1 # e.g. 0x00001111, where # of 1s = bit_depth
128 | for c in range(3):
129 | raw_data[c::3] = (colors[c] & mask).reshape(raw_data.size // 3)
130 |
131 | return _flat_array_to_str(raw_data[header_data_size:], bits_per_index=bit_depth, end_str=end_str)
132 |
133 |
134 | def save_text_as_image_file(text_data: str, input_surface: pygame.Surface, filepath: str,
135 | bit_depth_range=(1, 4), end_str=chr(0), resize_mode='smooth'):
136 | to_save = write_text_to_surface(text_data, input_surface,
137 | bit_depth_range=bit_depth_range,
138 | end_str=end_str,
139 | resize_mode=resize_mode)
140 | pygame.image.save(to_save, filepath)
141 |
142 |
143 | def load_text_from_image_file(filepath: str, end_str=chr(0)) -> str:
144 | img = pygame.image.load(filepath)
145 | return read_text_from_surface(img, end_str=end_str)
146 |
147 |
148 | def _str_to_flat_array(data: str, bits_per_index=2):
149 | """
150 | Example input: "abc" with bits_per_index=2
151 | raw_bytes: [01100001 01100010 01100011] (aka 97 98 99)
152 | raw_bits: [1 0 0 0 0 1 1 0 0 1 0 0 0 1 1 0 1 1 0 0 0 1 1 0]
153 | res: [01 00 10 01 10 00 10 01 11 00 10 01]
154 | """
155 | raw_bytes = numpy.array(bytearray(data, 'utf-8'), dtype='uint8')
156 | raw_bits = numpy.array([0] * 8 * raw_bytes.size, dtype='uint8')
157 | for i in range(8):
158 | raw_bits[i:i + 8 * raw_bytes.size:8] = (raw_bytes & (1 << i)) >> i
159 | if raw_bits.size % bits_per_index > 0:
160 | # pad end with 0s if bits_per_index doesn't cleanly divide raw_bits
161 | raw_bits = numpy.pad(raw_bits, (0, bits_per_index - (raw_bits.size % bits_per_index)),
162 | 'constant', constant_values=0)
163 | res = numpy.array([0] * math.ceil(raw_bits.size / bits_per_index), dtype='uint8')
164 | for j in range(bits_per_index):
165 | res |= raw_bits[j::bits_per_index] << j
166 | return res
167 |
168 |
169 | def _flat_array_to_str(arr, bits_per_index=2, end_str=chr(0)) -> str:
170 | """Reverse of _str_to_flat_array"""
171 | raw_bits = numpy.array([0] * (arr.size * bits_per_index), dtype='uint8')
172 | for j in range(bits_per_index):
173 | raw_bits[j::bits_per_index] = (arr & (1 << j)) >> j
174 |
175 | overflow = raw_bits.size % 8
176 | if overflow > 0:
177 | raw_bits = numpy.resize(raw_bits, (raw_bits.size - overflow,))
178 |
179 | raw_bytes = numpy.array([0] * (raw_bits.size // 8), dtype='uint8')
180 | for i in range(8):
181 | raw_bytes |= raw_bits[i::8] << i
182 |
183 | as_bytes = raw_bytes.tobytes()
184 | if end_str.encode("utf-8") in as_bytes:
185 | return as_bytes[0:as_bytes.index(end_str.encode("utf-8"))].decode("utf-8")
186 | else:
187 | return as_bytes.decode("utf-8")
188 |
189 |
190 | # yoinked from https://stackoverflow.com/questions/12173774/how-to-modify-bits-in-an-integer
191 | def _set_bit(v, index, x) -> int:
192 | """Set the index:th bit of v to 1 if x is truthy, else to 0, and return the new value."""
193 | mask = 1 << index # Compute mask, an integer with just bit 'index' set.
194 | v &= ~mask # Clear the bit indicated by the mask (if x is False)
195 | if x:
196 | v |= mask # If x was True, set the bit indicated by the mask.
197 | return v # Return the result, we're done.
198 |
199 |
200 | if __name__ == "__main__":
201 | input_filename = "data/splash.png"
202 | output_filename = "data/splash_output.png"
203 | _end_str = "~END~"
204 | img = pygame.image.load(input_filename)
205 |
206 | import json
207 | with open("data/arrival.json") as f:
208 | input_data_as_json = json.load(f)
209 |
210 | input_data = json.dumps(input_data_as_json, ensure_ascii=True)
211 | input_data = input_data * 10
212 |
213 | new_surf = write_text_to_surface(input_data, img, bit_depth_range=(1, 5), end_str=_end_str, resize_mode='integer')
214 | pygame.image.save(new_surf, output_filename)
215 |
216 | output_data_nosave = read_text_from_surface(new_surf, end_str=_end_str)
217 | output_data_from_img = load_text_from_image_file(output_filename, end_str=_end_str)
218 |
219 | print("input_data:", input_data)
220 | print("output_data_nosave:", output_data_nosave)
221 | print("output_data_from_img:", output_data_from_img)
222 | print(f"input_data == output_data_nosave = {input_data == output_data_nosave}")
223 | print(f"input_data == output_data_from_img = {input_data == output_data_from_img}")
224 |
--------------------------------------------------------------------------------
/threedee.py:
--------------------------------------------------------------------------------
1 | from typing import List, Iterable
2 |
3 | import numpy
4 | import pygame
5 |
6 | from pygame import Vector3, Vector2
7 | import math
8 |
9 |
10 | def ortho_matrix(left, right, bottom, top, near_val, far_val):
11 | res = numpy.identity(4, dtype=numpy.float32)
12 | res.itemset((0, 0), float(2 / (right - left)))
13 | res.itemset((1, 1), float(2 / (top - bottom)))
14 | res.itemset((2, 2), float(-2 / (far_val - near_val)))
15 |
16 | t_x = -(right + left) / (right - left)
17 | t_y = -(top + bottom) / (top - bottom)
18 | t_z = -(far_val + near_val) / (far_val - near_val)
19 | res.itemset((0, 3), float(t_x))
20 | res.itemset((1, 3), float(t_y))
21 | res.itemset((2, 3), float(t_z))
22 |
23 | return res
24 |
25 |
26 | def perspective_matrix(fovy, aspect, z_near, z_far):
27 | f = 1 / math.tan(fovy / 2)
28 | res = numpy.identity(4, dtype=numpy.float32)
29 | res.itemset((0, 0), f / aspect)
30 | res.itemset((1, 1), f)
31 | res.itemset((2, 2), (z_far + z_near) / (z_near - z_far))
32 | res.itemset((3, 2), (2 * z_far * z_near) / (z_near - z_far))
33 | res.itemset((2, 3), -1)
34 | res.itemset((3, 3), 0)
35 | return res
36 |
37 |
38 | def get_matrix_looking_at(eye_xyz, target_xyz, up_vec):
39 | n = eye_xyz - target_xyz
40 | n.scale_to_length(1)
41 | u = up_vec.cross(n)
42 | v = n.cross(u)
43 | res = numpy.array([[u[0], u[1], u[2], (-u).dot(eye_xyz)],
44 | [v[0], v[1], v[2], (-v).dot(eye_xyz)],
45 | [n[0], n[1], n[2], (-n).dot(eye_xyz)],
46 | [0, 0, 0, 1]], dtype=numpy.float32)
47 | return res
48 |
49 |
50 | class Camera3D:
51 |
52 | def __init__(self):
53 | self.position = Vector3(0, 0, 0)
54 | self.direction: Vector3 = Vector3(0, 0, 1)
55 | self.up: Vector3 = Vector3(0, -1, 0)
56 | self.fov_degrees: float = 45 # vertical field of view
57 |
58 | def __repr__(self):
59 | return "{}(pos={}, dir={})".format(type(self).__name__, self.position, self.direction)
60 |
61 | def get_xform(self, surface_size):
62 | view_mat = get_matrix_looking_at(self.position, self.position + self.direction, self.up)
63 | proj_mat = perspective_matrix(self.fov_degrees / 180 * math.pi, surface_size[0] / surface_size[1], 0.5, 100000)
64 | return proj_mat @ view_mat
65 |
66 | def project_points_to_surface(self, screen_dims, points) -> List[Vector2]:
67 | camera_xform = self.get_xform(screen_dims)
68 | n = len(points)
69 |
70 | # coalesce all the points into a single numpy array
71 | point_list = numpy.ndarray((n, 4), dtype=numpy.float32)
72 | for i in range(n):
73 | pt = points[i]
74 | point_list[i] = (pt[0], pt[1], pt[2], 1)
75 |
76 | # transform the points through the camera's view matrix
77 | point_list = point_list.transpose()
78 | point_list = camera_xform.dot(point_list)
79 | point_list = point_list.transpose()
80 |
81 | res = []
82 | for i in range(n):
83 | w = point_list[i][3]
84 |
85 | if w > 0.001:
86 | x = screen_dims[0] * (0.5 + point_list[i][0] / w)
87 | y = screen_dims[1] * (0.5 + point_list[i][1] / w)
88 | res.append(Vector2(x, y))
89 | else:
90 | # means the point is behind the camera, and shouldn't be drawn
91 | res.append(None)
92 | return res
93 |
94 | def draw_line_3d(self, screen, p1: Vector3, p2: Vector3, color=(255, 255, 255), width=1):
95 | xformed_pts = self.project_points_to_surface(screen.get_size(), [p1, p2])
96 | if xformed_pts[0] is not None and xformed_pts[1] is not None:
97 | pygame.draw.line(screen, color, xformed_pts[0], xformed_pts[1], width=width)
98 |
99 | def draw_lines_3d(self, screen, lines):
100 | """
101 | lines: list of tuples (p1, p2, color, width)
102 | """
103 | all_pts = []
104 | for l in lines:
105 | all_pts.append(l[0])
106 | all_pts.append(l[1])
107 | all_xformed_pts = self.project_points_to_surface(screen.get_size(), all_pts)
108 | for i in range(0, len(all_xformed_pts) // 2):
109 | p1 = all_xformed_pts[i * 2]
110 | p2 = all_xformed_pts[i * 2 + 1]
111 | color = lines[i][2]
112 | width = lines[i][3]
113 | if p1 is not None and p2 is not None:
114 | pygame.draw.line(screen, color, p1, p2, width)
115 |
116 |
117 | def gen_cube(angle, size, center, color):
118 | res = []
119 | pts = []
120 | for x in (-1, 1):
121 | for z in (-1, 1):
122 | xz = Vector2(x, z)
123 | xz = xz.rotate(angle)
124 | for y in (-1, 1):
125 | pts.append(Vector3(xz[0], y, xz[1]) * (size / 2) + center)
126 |
127 | pt = pts[-1]
128 | for n in pts[:len(pts)-1]:
129 | if abs((pt - n).length() - size) <= 0.1:
130 | res.append((pt, n, color, 1))
131 | return res
132 |
133 |
134 | if __name__ == "__main__":
135 | # call it to see demo
136 | import sys
137 |
138 | pygame.init()
139 |
140 | screen = pygame.display.set_mode((600, 300), pygame.RESIZABLE)
141 |
142 | clock = pygame.time.Clock()
143 |
144 | camera = Camera3D()
145 | camera.position = Vector3(0, 10, -50)
146 |
147 | angle = 0
148 | lines = []
149 |
150 | import random
151 |
152 | cubes = []
153 | for _ in range(0, 10):
154 | angle = random.random() * 360
155 | speed = random.random() * 1
156 | size = 10 + random.random() * 30
157 | x = -100 + random.random() * 200
158 | z = 100 + random.random() * 40
159 | y = size / 2
160 | color = [random.randint(0, 255) for _ in range(3)]
161 | cubes.append([angle, speed, size, Vector3(x, y, z), color])
162 |
163 | while True:
164 | events = pygame.event.get()
165 | for e in events:
166 | if e.type == pygame.QUIT:
167 | sys.exit(0)
168 | if e.type == pygame.KEYDOWN:
169 | if e.key == pygame.K_ESCAPE:
170 | sys.exit(0)
171 | elif e.key == pygame.K_i:
172 | print("camera = " + str(camera))
173 |
174 | keys_held = pygame.key.get_pressed()
175 | if keys_held[pygame.K_LEFT] or keys_held[pygame.K_RIGHT]:
176 | xz = Vector2(camera.direction.x, camera.direction.z)
177 | xz = xz.rotate(1 * (1 if keys_held[pygame.K_LEFT] else -1))
178 | camera.direction.x = xz[0]
179 | camera.direction.z = xz[1]
180 | camera.direction.scale_to_length(1)
181 |
182 | if keys_held[pygame.K_UP] or keys_held[pygame.K_DOWN]:
183 | camera.direction.y += 0.01 * (1 if keys_held[pygame.K_UP] else -1)
184 | camera.direction.scale_to_length(1)
185 |
186 | ms = 1
187 | xz = Vector2(camera.position.x, camera.position.z)
188 | view_xz = Vector2(camera.direction.x, camera.direction.z)
189 | view_xz.scale_to_length(1)
190 |
191 | if keys_held[pygame.K_a]:
192 | xz = xz + ms * view_xz.rotate(90)
193 | if keys_held[pygame.K_d]:
194 | xz = xz + ms * view_xz.rotate(-90)
195 | if keys_held[pygame.K_w]:
196 | xz = xz + ms * view_xz
197 | if keys_held[pygame.K_s]:
198 | xz = xz + ms * view_xz.rotate(180)
199 | camera.position.x = xz[0]
200 | camera.position.z = xz[1]
201 |
202 | screen.fill((0, 0, 0))
203 |
204 | lines = []
205 | for c in cubes:
206 | c[0] += c[1] # rotate
207 | lines.extend(gen_cube(c[0], c[2], c[3], c[4]))
208 |
209 | camera.draw_lines_3d(screen, lines)
210 |
211 | pygame.display.update()
212 | pygame.display.set_caption(str(int(clock.get_fps())))
213 | clock.tick(60)
214 |
--------------------------------------------------------------------------------
/video.py:
--------------------------------------------------------------------------------
1 | import pygame
2 | import numpy
3 | import cv2
4 |
5 | import math
6 | import time
7 | import os
8 | import errno
9 |
10 | class Video:
11 | """A video playback utility for pygame.
12 |
13 | Requires:
14 | pygame (pip install pygame)
15 | numpy (pip install numpy)
16 | cv2 (pip install opencv-python)
17 |
18 | This class streams the image data from a video file and provides it as a pygame Surface. There are methods that
19 | control video playback, such as play, pause, and set_frame, as well as methods that provide information about the
20 | video itself, like the dimensions and the FPS.
21 |
22 | The video playback tries to play in "real time", meaning that if you call play, and then call get_surface 5 seconds
23 | later, you'll get the video frame at the 5 second mark.
24 |
25 | However, this class is also lazy, meaning that video data is only processed when get_surface is called. There is
26 | no asynchronous background processing, and it doesn't cache any video data in memory besides the current frame.
27 | You don't need to advance frames manually or update the video each frame to keep it "going".
28 |
29 | This class does not provide sound from video files, only the images.
30 |
31 | Example Usage:
32 |
33 | import pygame
34 | import video
35 |
36 | pygame.init()
37 |
38 | screen = pygame.display.set_mode((1080, 720))
39 | clock = pygame.time.Clock()
40 |
41 | vid = video.Video("your_file_goes_here.mp4")
42 | vid.play()
43 |
44 | while True:
45 | for e in pygame.event.get():
46 | if e.type == pygame.QUIT:
47 | raise SystemExit
48 | screen.fill((0, 0, 0))
49 | screen.blit(vid.get_surface(), (0, 0))
50 |
51 | pygame.display.flip()
52 | clock.tick(60)
53 | """
54 |
55 | class _PlaybackState:
56 | def __init__(self, start_frame, playing, t):
57 | self.frame = start_frame
58 | self.playing = playing
59 | self.t = t
60 |
61 | def __init__(self, filename, fps=0):
62 | """Inits a new Video.
63 |
64 | The video will be paused initially.
65 |
66 | filename: The path to the video file.
67 | fps: The playback framerate for the video. If 0, the native FPS will be used.
68 | """
69 | if not os.path.isfile(filename):
70 | raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), filename)
71 |
72 | self._filename = filename
73 | self._vid = cv2.VideoCapture(filename)
74 | self._vid_frame = 0 # the frame the next _vid.read() call will give.
75 |
76 | self._frame_count = int(self._vid.get(cv2.CAP_PROP_FRAME_COUNT))
77 | self._frame_width = int(self._vid.get(cv2.CAP_PROP_FRAME_WIDTH))
78 | self._frame_height = int(self._vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
79 | self._fps = int(self._vid.get(cv2.CAP_PROP_FPS)) if fps <= 0 else fps
80 |
81 | self._buf = pygame.Surface((self._frame_width, self._frame_height))
82 | self._cur_frame = -1 # the frame that was last drawn to _buf
83 |
84 | self._playback_state = Video._PlaybackState(0, False, time.time())
85 | self._final_frame = float('inf')
86 |
87 | def __repr__(self):
88 | return f"{type(self).__name__}({self._filename})"
89 |
90 | def play(self, loops=0, maxtime=0):
91 | """Begins video playback.
92 |
93 | loops: How many times playback should loop. If 0, it will repeat forever.
94 | maxtime: How long in seconds the video should play. If 0, it will play forever.
95 | """
96 | if self.is_paused():
97 | self._playback_state.playing = True
98 | self._playback_state.t = time.time()
99 | self.frame = self._playback_state.frame % self.get_frame_count() # go back to first loop
100 | self._final_frame = self._calc_final_frame(loops, maxtime)
101 |
102 | def set_frame(self, n):
103 | """Jumps to a specific frame."""
104 | self._playback_state.frame = n
105 | self._playback_state.t = time.time()
106 |
107 | def pause(self):
108 | """Pauses the video."""
109 | if not self.is_paused():
110 | self._playback_state.frame = self.get_current_frame(wrapped=False)
111 | self._playback_state.playing = False
112 |
113 | def is_paused(self):
114 | """Whether the video is paused.
115 |
116 | Note that if a video has finished (e.g. finished looping), it will still be considered unpaused.
117 | """
118 | return not self._playback_state.playing
119 |
120 | def is_finished(self) -> bool:
121 | """Whether the termination condition passed into play has been reached."""
122 | return self.get_current_frame() >= self._final_frame
123 |
124 | def get_surface(self) -> pygame.Surface:
125 | """Returns the video's image data for the current frame.
126 |
127 | This is where the bulk of this class's work is performed. This method calculates the current frame (based on the
128 | current time and other factors), and if the frame has changed since the last call to this method, it reads
129 | video data from the file and blits it onto a surface, which is returned.
130 |
131 | If the buffer surface is already up-to-date, this method returns it instantly.
132 | If the video has ended, (i.e. the termination condition passed into play() has been reached), a blank surface is returned.
133 | """
134 | cur_frame = self.get_current_frame(wrapped=False)
135 | if cur_frame >= self._final_frame:
136 | self._buf.fill((0, 0, 0)) # video is over, you get a black screen
137 | return self._buf
138 | else:
139 | self._draw_frame_to_surface_if_necessary(cur_frame)
140 | return self._buf
141 |
142 | def _draw_frame_to_surface_if_necessary(self, frame_n):
143 | frame_n %= self._frame_count
144 | if self._cur_frame == frame_n:
145 | return # this frame is already drawn
146 |
147 | if self._vid_frame > frame_n:
148 | # we have to loop back to the beginning
149 | self._vid_frame = 0
150 | self._vid = cv2.VideoCapture(self._filename)
151 |
152 | success, frame_data = None, None
153 | for _ in range(frame_n - self._vid_frame + 1):
154 | success, next_frame_data = self._vid.read()
155 | if not success:
156 | # sometimes CAP_PROP_FRAME_COUNT will be straight-up wrong, indicating more or less frames
157 | # than there actually are. In that case we just... skip or freeze the final frames, I guess?
158 | continue
159 | else:
160 | frame_data = next_frame_data
161 | self._vid_frame += 1
162 |
163 | if frame_data is not None:
164 | pygame.surfarray.blit_array(self._buf, numpy.flip(numpy.rot90(frame_data[::-1])))
165 | self._cur_frame = self._vid_frame - 1
166 |
167 | def get_width(self) -> int:
168 | """The width of the video in pixels."""
169 | return self._frame_width
170 |
171 | def get_height(self) -> int:
172 | """The height of the video in pixels."""
173 | return self._frame_height
174 |
175 | def get_size(self) -> (int, int):
176 | """The dimensions of the video in pixels."""
177 | return self.get_width(), self.get_height()
178 |
179 | def get_current_frame(self, wrapped=True) -> int:
180 | """The current frame number.
181 |
182 | wrapped: if True, the result will be less than frame_count.
183 | if False, the result may be greater or equal to frame_count, indicating the video has looped.
184 | """
185 | if self.is_paused():
186 | return self._playback_state.frame
187 | else:
188 | cur_time = time.time()
189 | start_time = self._playback_state.t
190 | elapsed_frames = int(self.get_fps() * (cur_time - start_time))
191 | cur_frame = self._playback_state.frame + elapsed_frames
192 |
193 | return cur_frame % self._frame_count if wrapped else cur_frame
194 |
195 | def get_fps(self) -> int:
196 | """The frames per second at which the video will play."""
197 | return self._fps
198 |
199 | def get_frame_count(self) -> int:
200 | """The number of frames in the video."""
201 | return self._frame_count
202 |
203 | def get_duration(self) -> float:
204 | """The duration of the video in seconds."""
205 | return self._frame_count / self._fps
206 |
207 | def _calc_final_frame(self, loops, maxtime):
208 | res = float('inf')
209 | if loops > 0:
210 | res = min(res, loops * self.get_frame_count())
211 | if maxtime > 0:
212 | res = min(res, math.ceil(self.get_fps() * maxtime))
213 | return res
214 |
--------------------------------------------------------------------------------
/warp.py:
--------------------------------------------------------------------------------
1 | import typing
2 |
3 | import pygame
4 | import numpy
5 | import cv2
6 |
7 |
8 | def warp(surf: pygame.Surface,
9 | warp_pts,
10 | smooth=True,
11 | out: pygame.Surface = None) -> typing.Tuple[pygame.Surface, pygame.Rect]:
12 | """Stretches a pygame surface to fill a quad using cv2's perspective warp.
13 |
14 | Args:
15 | surf: The surface to transform.
16 | warp_pts: A list of four xy coordinates representing the polygon to fill.
17 | Points should be specified in clockwise order starting from the top left.
18 | smooth: Whether to use linear interpolation for the image transformation.
19 | If false, nearest neighbor will be used.
20 | out: An optional surface to use for the final output. If None or not
21 | the correct size, a new surface will be made instead.
22 |
23 | Returns:
24 | [0]: A Surface containing the warped image.
25 | [1]: A Rect describing where to blit the output surface to make its coordinates
26 | match the input coordinates.
27 | """
28 | if len(warp_pts) != 4:
29 | raise ValueError("warp_pts must contain four points")
30 |
31 | w, h = surf.get_size()
32 | is_alpha = surf.get_flags() & pygame.SRCALPHA
33 |
34 | # XXX throughout this method we need to swap x and y coordinates
35 | # when we pass stuff between pygame and cv2. I'm not sure why .-.
36 | src_corners = numpy.float32([(0, 0), (0, w), (h, w), (h, 0)])
37 | quad = [tuple(reversed(p)) for p in warp_pts]
38 |
39 | # find the bounding box of warp points
40 | # (this gives the size and position of the final output surface).
41 | min_x, max_x = float('inf'), -float('inf')
42 | min_y, max_y = float('inf'), -float('inf')
43 | for p in quad:
44 | min_x, max_x = min(min_x, p[0]), max(max_x, p[0])
45 | min_y, max_y = min(min_y, p[1]), max(max_y, p[1])
46 | warp_bounding_box = pygame.Rect(int(min_x), int(min_y),
47 | int(max_x - min_x),
48 | int(max_y - min_y))
49 |
50 | shifted_quad = [(p[0] - min_x, p[1] - min_y) for p in quad]
51 | dst_corners = numpy.float32(shifted_quad)
52 |
53 | mat = cv2.getPerspectiveTransform(src_corners, dst_corners)
54 |
55 | orig_rgb = pygame.surfarray.pixels3d(surf)
56 |
57 | flags = cv2.INTER_LINEAR if smooth else cv2.INTER_NEAREST
58 | out_rgb = cv2.warpPerspective(orig_rgb, mat, warp_bounding_box.size, flags=flags)
59 |
60 | if out is None or out.get_size() != out_rgb.shape[0:2]:
61 | out = pygame.Surface(out_rgb.shape[0:2], pygame.SRCALPHA if is_alpha else 0)
62 |
63 | pygame.surfarray.blit_array(out, out_rgb)
64 |
65 | if is_alpha:
66 | orig_alpha = pygame.surfarray.pixels_alpha(surf)
67 | out_alpha = cv2.warpPerspective(orig_alpha, mat, warp_bounding_box.size, flags=flags)
68 | alpha_px = pygame.surfarray.pixels_alpha(out)
69 | alpha_px[:] = out_alpha
70 | else:
71 | out.set_colorkey(surf.get_colorkey())
72 |
73 | # XXX swap x and y once again...
74 | return out, pygame.Rect(warp_bounding_box.y, warp_bounding_box.x,
75 | warp_bounding_box.h, warp_bounding_box.w)
76 |
77 |
78 | if __name__ == "__main__":
79 | pygame.init()
80 |
81 | screen = pygame.display.set_mode((640, 480))
82 | pygame.display.set_caption("warp.py")
83 |
84 | frog_img = pygame.image.load("data/frog.png").convert_alpha()
85 | frog_img = pygame.transform.scale(frog_img, (frog_img.get_width() * 5,
86 | frog_img.get_height() * 5))
87 | default_rect = frog_img.get_rect(center=screen.get_rect().center)
88 | warped_frog_img = None
89 |
90 | corners = [default_rect.topleft, default_rect.topright,
91 | default_rect.bottomright, default_rect.bottomleft]
92 | held_corner_idx = -1
93 |
94 | automatic_demo_mode = True
95 | t = 0
96 |
97 | clock = pygame.time.Clock()
98 | font = pygame.font.SysFont("", 24)
99 |
100 | while True:
101 | for e in pygame.event.get():
102 | if e.type == pygame.QUIT:
103 | raise SystemExit()
104 | elif e.type == pygame.KEYDOWN:
105 | if e.key == pygame.K_ESCAPE:
106 | raise SystemExit()
107 | elif e.key == pygame.K_RETURN:
108 | # [Enter] = toggle demo mode
109 | automatic_demo_mode = not automatic_demo_mode
110 | elif e.key == pygame.K_r:
111 | # [R] = reset corners
112 | corners = [default_rect.topleft, default_rect.topright,
113 | default_rect.bottomright, default_rect.bottomleft]
114 | elif e.type == pygame.MOUSEBUTTONDOWN:
115 | if e.button == 1:
116 | # [LMB] = move a corner to a new position and start dragging it
117 | automatic_demo_mode = False # enter manual mode
118 |
119 | # find the point closest to the click
120 | epos_vector = pygame.Vector2(e.pos)
121 | best_idx, best_dist = 0, float('inf')
122 | for idx, c in enumerate(corners):
123 | c_dist = epos_vector.distance_to(c)
124 | if c_dist < best_dist:
125 | best_idx = idx
126 | best_dist = c_dist
127 | held_corner_idx = best_idx # indicate we're dragging that point
128 | corners[best_idx] = e.pos # move the point to the click location
129 | elif e.type == pygame.MOUSEBUTTONUP:
130 | if e.button == 1:
131 | held_corner_idx = -1 # release the point we're dragging
132 |
133 | keys = pygame.key.get_pressed()
134 |
135 | # move the 'held point' to the mouse's location
136 | mouse_pos = pygame.mouse.get_pos()
137 | if mouse_pos is not None and held_corner_idx >= 0:
138 | corners[held_corner_idx] = mouse_pos
139 |
140 | # make the points oscillate in circles if we're in 'demo mode'
141 | if automatic_demo_mode:
142 | perturbs = [pygame.Vector2() for _ in range(4)]
143 | for idx, pert in enumerate(perturbs):
144 | pert.from_polar(((idx + 1) * 15, (5 - idx) * 30 * t)) # circular motion
145 | pts_to_use = [pert + pygame.Vector2(c) for c, pert in zip(corners, perturbs)]
146 | else:
147 | pts_to_use = corners
148 |
149 | screen.fill((40, 45, 50))
150 |
151 | # generate the warped image
152 | warped_frog_img, warped_pos = warp(
153 | frog_img,
154 | pts_to_use,
155 | smooth=not keys[pygame.K_SPACE], # toggle smoothing while [Space] is held
156 | out=warped_frog_img)
157 |
158 | # draw green border around the warped image
159 | pygame.draw.rect(
160 | screen, "limegreen",
161 | warped_frog_img.get_rect(topleft=warped_pos.topleft), width=1)
162 | border_text = font.render(f"pos=({warped_pos.x}, {warped_pos.y})", True, "lime")
163 | screen.blit(border_text, (warped_pos.x, warped_pos.y - border_text.get_height() - 2))
164 |
165 | # draw red warp guidelines
166 | pygame.draw.line(screen, "red2", pts_to_use[0], pts_to_use[2], width=1)
167 | pygame.draw.line(screen, "red2", pts_to_use[1], pts_to_use[3], width=1)
168 | for i in range(len(pts_to_use)):
169 | pygame.draw.line(screen, "red2", pts_to_use[i], pts_to_use[(i + 1) % 4], width=2)
170 |
171 | # draw labels on warp points
172 | for i, pt in enumerate(pts_to_use):
173 | text_img = font.render(f"p{i}=({int(pt[0])}, {int(pt[1])})", True, "red")
174 | scr_pos = (pt[0] + 2, pt[1]) if i in (1, 2) else (pt[0] - text_img.get_width() - 2, pt[1])
175 | screen.blit(text_img, scr_pos)
176 |
177 | # blit actual warped image
178 | screen.blit(warped_frog_img, warped_pos)
179 |
180 | pygame.display.flip()
181 | t += clock.tick(60) / 1000.0
182 |
--------------------------------------------------------------------------------
/worms.py:
--------------------------------------------------------------------------------
1 | import pygame
2 |
3 | import math
4 | import random
5 |
6 | # "Infinite Worms" Screensaver Project by Ghast
7 |
8 | WINDOW_TITLE = "Infinite Worms"
9 | N_AGENTS = 48
10 | GRID_DIMS = 64, 32
11 | UPSCALE = 8
12 | SCREEN_SIZE = 640, 320
13 |
14 | WORM_SPEED = 30 # moves per sec
15 | COLOR_OFFSET_VARIATION = 3 # secs
16 | FPS = 60
17 |
18 | SHRINK_RATE = 4
19 | FADE_RATE = 12
20 | BLUR_RADIUS = 1
21 |
22 | HUE_SHIFT_RATE = 12 # degrees per sec
23 |
24 | SATURATION_RANGE = (0.666, 1.0)
25 | SATURATION_SHIFT_RATE = 60 # seconds per cycle
26 |
27 | VALUE_SHIFT_RANGE = (0.666, 1.0)
28 | VALUE_SHIFT_RATE = 16 # seconds per cycle
29 |
30 |
31 | class Grid:
32 |
33 | def __init__(self, dims, agents):
34 | self.t = 0
35 | self.ticks = 0
36 | self.dims = dims
37 | self.surf = pygame.Surface(dims).convert_alpha()
38 |
39 | self.pheromones = pygame.Surface(dims)
40 | self.phero_fade = pygame.Surface(dims)
41 | self.phero_fade.fill("black")
42 | self.phero_fade.set_alpha(1)
43 |
44 | self.upscale = pygame.Surface((dims[0] * UPSCALE, dims[1] * UPSCALE))
45 |
46 | self.fade = pygame.Surface((dims[0] * UPSCALE, dims[1] * UPSCALE))
47 | self.fade.fill("black")
48 | self.fade.set_alpha(1)
49 |
50 | self.agents = agents
51 |
52 | def update(self, dt):
53 | self.surf.fill((0, 0, 0, 0))
54 | for a in self.agents:
55 | a.update(self, dt)
56 | self.pheromones.blit(self.phero_fade, (0, 0))
57 |
58 | if self.ticks % SHRINK_RATE == 0:
59 | shrink = pygame.transform.smoothscale(self.upscale, (self.upscale.get_width() - 4, self.upscale.get_height() - 2))
60 | self.upscale.fill("black")
61 | self.upscale.blit(shrink, (2, 1))
62 |
63 | if self.ticks % FADE_RATE == 0:
64 | self.upscale.blit(self.fade, (0, 0))
65 | self.upscale = pygame.transform.gaussian_blur(self.upscale, BLUR_RADIUS)
66 |
67 | ups = pygame.transform.scale(self.surf, self.upscale.get_size())
68 | self.upscale.blit(ups, (0, 0))
69 |
70 | self.t += dt
71 | self.ticks += 1
72 |
73 | def render(self, screen: pygame.Surface, mode='normal'):
74 | if mode == 'normal':
75 | screen.blit(pygame.transform.smoothscale(self.upscale, screen.get_size()), (0, 0))
76 | elif mode == 'pheromone':
77 | screen.blit(pygame.transform.scale(self.pheromones, screen.get_size()), (0, 0))
78 | else:
79 | screen.fill("black")
80 |
81 |
82 | class Agent:
83 |
84 | def __init__(self, xy, period=1 / WORM_SPEED, offs=0.0):
85 | self.color = pygame.Color("red")
86 | self.hsva = self.color.hsva
87 | self.xy = xy
88 | self.period = period
89 | self.cooldown = period
90 | self.t = 0
91 | self.offs = offs
92 |
93 | def update(self, grid, dt):
94 | self.cooldown -= dt
95 | if self.cooldown <= 0:
96 | self.cooldown = self.period
97 | self.xy = self.calc_next_pos(grid)
98 | grid.pheromones.set_at(self.xy, self.color)
99 | grid.surf.set_at(self.xy, self.color)
100 |
101 | self.t += dt
102 | t = self.t + self.offs
103 |
104 | self.hsva = ((t * HUE_SHIFT_RATE) % 360,
105 | SATURATION_RANGE[0] + (SATURATION_RANGE[1] - SATURATION_RANGE[0]) * (0.5 + math.cos(2 * math.pi * t / SATURATION_SHIFT_RATE) / 2),
106 | VALUE_SHIFT_RANGE[0] + (VALUE_SHIFT_RANGE[1] - VALUE_SHIFT_RANGE[0]) * (0.5 + math.cos(2 * math.pi * t / VALUE_SHIFT_RATE) / 2),
107 | 1)
108 | self.color.hsva = int(self.hsva[0]), int(100 * self.hsva[1]), int(100 * self.hsva[2]), int(100 * self.hsva[3])
109 |
110 | def _evaluate(self, xy, grid):
111 | if xy[0] < 0 or xy[1] < 0 or xy[0] >= grid.dims[0] or xy[1] >= grid.dims[1]:
112 | return 1000 # don't wrap
113 | else:
114 | color = pygame.Color(grid.pheromones.get_at(xy))
115 | return color.hsva[2]
116 |
117 | def calc_next_pos(self, grid):
118 | neighbors = [(-1, 0), (1, 0), (0, -1), (0, 1)]
119 | random.shuffle(neighbors)
120 | xys = [(self.xy[0] + n[0], self.xy[1] + n[1]) for n in neighbors]
121 | xys.sort(key=lambda n: self._evaluate(n, grid))
122 | return xys[0]
123 |
124 |
125 | def main():
126 | pygame.init()
127 |
128 | screen = pygame.display.set_mode(SCREEN_SIZE, pygame.RESIZABLE)
129 | pygame.display.set_caption(f"{WINDOW_TITLE}")
130 | clock = pygame.time.Clock()
131 | dt = 0
132 |
133 | agents = [Agent((random.randint(0, GRID_DIMS[0] - 1), random.randint(0, GRID_DIMS[1] - 1)),
134 | offs=i / N_AGENTS * COLOR_OFFSET_VARIATION) for i in range(N_AGENTS)]
135 | grid = Grid(GRID_DIMS, agents)
136 |
137 | running = True
138 | while running:
139 | for e in pygame.event.get():
140 | if e.type == pygame.QUIT or (e.type == pygame.KEYDOWN and e.key == pygame.K_ESCAPE):
141 | running = False
142 | if e.type == pygame.KEYDOWN:
143 | if e.key == pygame.K_F4:
144 | pygame.display.toggle_fullscreen()
145 |
146 | keys = pygame.key.get_pressed()
147 | mode = 'normal' if not keys[pygame.K_SPACE] else 'pheromone'
148 |
149 | grid.update(dt)
150 |
151 | screen.fill("black")
152 | grid.render(screen, mode=mode)
153 |
154 | pygame.display.flip()
155 | dt = clock.tick(FPS) / 1000
156 |
157 | if grid.ticks % 15 == 14:
158 | fps = clock.get_fps()
159 | pygame.display.set_caption(f"{WINDOW_TITLE} [FPS={fps:.1f}]")
160 |
161 | main()
162 |
163 |
--------------------------------------------------------------------------------