├── .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 | ![life_demo.gif](screenshots/life_demo.gif?raw=true "Life Demo") 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 | ![colorswirl.gif](screenshots/colorswirl.gif?raw=true "Colorswirl Demo") 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 | ![fractal_zoom.gif](screenshots/fractal_zoom.gif?raw=true "Fractal Demo") 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 | ![rainbowize.gif](screenshots/rainbowize.gif?raw=true "Rainbowize Demo") 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 | ![warp.gif](screenshots/warp.gif?raw=true "Warp Demo") 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 | ![lut.gif](screenshots/lut.gif?raw=true "LUT Demo") 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 | ![worms.gif](screenshots/worms.gif?raw=true "Infinite Worms Demo") 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 | ![video_demo.gif](screenshots/video_demo.gif?raw=true "Video Demo") 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 | ![stega_demo.png](screenshots/stega_demo.png?raw=true "Steganography Demo") 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 | ![gifreader_demo.gif](screenshots/gifreader_demo.gif?raw=true "gifreader_demo") 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 | ![benchmark_results.png](screenshots/benchmark_results.png?raw=true "benchmark_results") 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 | ![n=1000animated.gif](screenshots/n=1000animated.gif?raw=true "n=1000 animated") 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 | ![n=1000.PNG](screenshots/n=1000.PNG?raw=true "n=1000") 159 | 160 | ### N=10,000 161 | ![n=10000.PNG](screenshots/n=10000.PNG?raw=true "n=10000") 162 | 163 | ### N=100,000 164 | ![n=100000.PNG](screenshots/n=100000.PNG?raw=true "n=100000") 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 | --------------------------------------------------------------------------------