├── .gitignore ├── __init__.py ├── auutil.py ├── color_conversion.py ├── crop.py ├── draw_shapes.py ├── ffmpeg.py ├── geometry.py ├── histogram_equalize.py ├── imutil.py ├── imwarp.py ├── itertools.py ├── license.md ├── list_files.py ├── memoize.py ├── mosaic.py ├── multiworker.py ├── pixelated.py ├── plot_images.py ├── progress.py ├── rainbow.py ├── readme.md ├── setup.py ├── specshow.py ├── stratified.py ├── timing.py └── wavshow.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | .ipynb_checkpoints 3 | *.pyc 4 | .vscode -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kylemcdonald/python-utils/e89c4967b57b0d0a41a7b712e2bbf47a1ecd393b/__init__.py -------------------------------------------------------------------------------- /auutil.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def autrim(audio, thresh=1e-4): 4 | start = np.argmin(audio < thresh) 5 | end = np.argmax(audio[::-1] > thresh) 6 | return audio[start:-end-1], start, end -------------------------------------------------------------------------------- /color_conversion.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | def to_single_rgb(img): 3 | img = np.asarray(img) 4 | if len(img.shape) == 4: # take first frame from animations 5 | return img[0,:,:,:] 6 | if len(img.shape) == 2: # convert gray to rgb 7 | img = img[:,:,np.newaxis] 8 | return np.repeat(img, 3, 2) # might np.tile(img, [1,1,3]) be faster? 9 | if img.shape[-1] == 4: # drop alpha 10 | return img[:,:,:3] 11 | else: 12 | return img 13 | 14 | def to_single_gray(img): 15 | img = np.asarray(img) 16 | if len(img.shape) == 2: 17 | return img[:,:,np.newaxis] 18 | elif img.shape[2] == 3: 19 | return img.mean(axis=2) 20 | else: 21 | return img 22 | 23 | def rb_swap(img): 24 | if len(img.shape) < 3: 25 | return img 26 | if img.shape[2] == 3: 27 | return img[...,(2,1,0)] 28 | if img.shape[2] == 4: 29 | return img[...,(2,1,0,3)] -------------------------------------------------------------------------------- /crop.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def safe_crop(arr, tblr, fill=None): 4 | n,s,w,e = tblr 5 | shape = np.asarray(arr.shape) 6 | shape[:2] = s - n, e - w 7 | no, so, wo, eo = 0, shape[0], 0, shape[1] 8 | if n < 0: 9 | no += -n 10 | n = 0 11 | if w < 0: 12 | wo += -w 13 | w = 0 14 | if s >= arr.shape[0]: 15 | so -= s - arr.shape[0] 16 | s = arr.shape[0] 17 | if e >= arr.shape[1]: 18 | eo -= e - arr.shape[1] 19 | e = arr.shape[1] 20 | cropped = arr[n:s,w:e] 21 | if fill is None: 22 | return cropped 23 | out = np.empty(shape, dtype=arr.dtype) 24 | out.fill(fill) 25 | try: 26 | out[no:so,wo:eo] = cropped 27 | except ValueError: 28 | # this happens when there is no overlap 29 | pass 30 | return out 31 | 32 | def place_inside(img, xy, output_shape, fill=0): 33 | ix,iy = xy 34 | ih,iw,ic = img.shape 35 | oh,ow = output_shape 36 | out = np.empty((oh,ow,ic), dtype=img.dtype) 37 | out.fill(fill) 38 | ox = 0 39 | oy = 0 40 | if ih < oh and iy == 0: 41 | oy = oh - ih 42 | if iw < ow and ix == 0: 43 | ox = ow - iw 44 | out[oy:oy+ih,ox:ox+iw] = img 45 | return out 46 | 47 | def inner_square_crop(img): 48 | size = np.asarray(img.shape[:2]) 49 | min_side = min(size) 50 | corner = (size - min_side) // 2 51 | return img[corner[0]:corner[0]+min_side, corner[1]:corner[1]+min_side] 52 | 53 | def outer_square_crop(img, fill=0): 54 | size = np.asarray(img.shape[:2]) 55 | max_side = max(size) 56 | output_shape = np.copy(img.shape) 57 | output_shape[:2] = max_side 58 | out = np.empty(output_shape, dtype=img.dtype) 59 | out.fill(fill) 60 | corner = (max_side - size) // 2 61 | out[corner[0]:corner[0]+size[0], corner[1]:corner[1]+size[1]] = img 62 | return out 63 | -------------------------------------------------------------------------------- /draw_shapes.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | 3 | def draw_line(canvas, pt1, pt2, r=1, stroke=None): 4 | pt1 = tuple(map(int, pt1)) 5 | pt2 = tuple(map(int, pt2)) 6 | cv2.line(canvas, pt1, pt2, stroke, thickness=r, lineType=cv2.LINE_AA) 7 | 8 | def draw_text(canvas, text, xy, color=0, scale=1, thickness=1, highlight=None, 9 | font_face=cv2.FONT_HERSHEY_SIMPLEX, antialias=False): 10 | l,t = tuple(map(int, xy)) 11 | (tw,th), baseline = cv2.getTextSize(text, font_face, scale, thickness) 12 | t += th + baseline - 1 13 | if highlight is not None: 14 | canvas[t-th-baseline-1:t,l:l+tw] = highlight 15 | cv2.putText(canvas, text, (l,t-baseline), font_face, scale, color, thickness, cv2.LINE_AA if antialias else 0) 16 | 17 | # for some reason a fill of 0 doesn't work, but 0.1 does work 18 | def draw_circle(canvas, xy, r=1, stroke=None, fill=None, thickness=1, antialias=False): 19 | x,y = tuple(map(int, xy)) 20 | line_type = cv2.LINE_AA if antialias else cv2.LINE_8 21 | if fill is not None: 22 | cv2.circle(canvas, (x,y), r, fill, -1, line_type) 23 | if stroke is not None: 24 | cv2.circle(canvas, (x,y), r, stroke, thickness, line_type) 25 | 26 | # @njit 27 | # def draw_circle(canvas, xy, r, fill): 28 | # x,y = xy 29 | # r2 = r * r 30 | # for i in range(canvas.shape[0]): 31 | # cy = i - y 32 | # cy2 = cy * cy 33 | # for j in range(canvas.shape[1]): 34 | # cx = j - x 35 | # ls = cx * cx + cy2 36 | # if ls < r2: 37 | # canvas[i,j] = fill 38 | 39 | def draw_rectangle_thin(canvas, tblr, fill=None, stroke=None): 40 | t,b,l,r = tblr 41 | ye = canvas.shape[0] - 1 42 | xe = canvas.shape[1] - 1 43 | t = int(min(max(t,0),ye)) 44 | b = int(min(max(b,0),ye)) 45 | l = int(min(max(l,0),xe)) 46 | r = int(min(max(r,0),xe)) 47 | if fill is not None: 48 | canvas[t:b,l:r] = fill 49 | if stroke is not None: 50 | b = int(max(b-1,0)) 51 | r = int(max(r-1,0)) 52 | try: 53 | canvas[t:b,l] = stroke 54 | except IndexError: 55 | pass 56 | try: 57 | canvas[t:b,r] = stroke 58 | except IndexError: 59 | pass 60 | try: 61 | canvas[t,l:r] = stroke 62 | except IndexError: 63 | pass 64 | try: 65 | r = int(min(r+1,xe)) 66 | canvas[b,l:r] = stroke 67 | except IndexError: 68 | pass 69 | 70 | def draw_rectangle(canvas, tblr, fill=None, stroke=None, thickness=1): 71 | draw_rectangle_thin(canvas, tblr, fill, stroke) 72 | for i in range(1, thickness): 73 | t,b,l,r = tblr 74 | draw_rectangle_thin(canvas, (t-i,b+i,l-i,r+i), fill, stroke) 75 | draw_rectangle_thin(canvas, (t+i,b-i,l+i,r-i), fill, stroke) 76 | 77 | def draw_rectangle_dlib(canvas, det, fill=None, stroke=None): 78 | rect = (det.top(), det.bottom(), det.left(), det.right()) 79 | draw_rectangle(canvas, rect, fill=fill, stroke=stroke) 80 | 81 | -------------------------------------------------------------------------------- /ffmpeg.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import subprocess as sp 3 | import os 4 | import time 5 | import ffmpeg 6 | DEVNULL = open(os.devnull, 'w') 7 | 8 | # attempts to handle all float/integer conversions with and without normalizing 9 | def convert_bit_depth(y, in_type, out_type, normalize=False): 10 | in_type = np.dtype(in_type).type 11 | out_type = np.dtype(out_type).type 12 | 13 | if normalize: 14 | peak = np.abs(y).max() 15 | if peak == 0: 16 | normalize = False 17 | 18 | if issubclass(in_type, np.floating): 19 | if normalize: 20 | y /= peak 21 | if issubclass(out_type, np.integer): 22 | y *= np.iinfo(out_type).max 23 | y = y.astype(out_type) 24 | elif issubclass(in_type, np.integer): 25 | if issubclass(out_type, np.floating): 26 | y = y.astype(out_type) 27 | if normalize: 28 | y /= peak 29 | elif issubclass(out_type, np.integer): 30 | in_max = peak if normalize else np.iinfo(in_type).max 31 | out_max = np.iinfo(out_type).max 32 | if out_max > in_max: 33 | y = y.astype(out_type) 34 | y *= (out_max / in_max) 35 | elif out_max < in_max: 36 | y /= (in_max / out_max) 37 | y = y.astype(out_type) 38 | return y 39 | 40 | def aureadmeta(fn): 41 | if not os.path.exists(fn): 42 | raise FileNotFoundError 43 | probe = ffmpeg.probe(fn) 44 | for stream in probe['streams']: 45 | if stream['codec_type'] == 'audio': 46 | meta = { 47 | 'channels': stream['channels'], 48 | 'sample_rate': int(stream['sample_rate']), 49 | 'duration': float(probe['format']['duration']) 50 | } 51 | return meta 52 | return None 53 | 54 | # auread should be combined with aureadmeta to not force the samplerate or input type if they are None 55 | def auread(filename, sr=44100, mono=False, normalize=True, in_type=np.int16, out_type=np.float32): 56 | in_type = np.dtype(in_type).type 57 | out_type = np.dtype(out_type).type 58 | channels = 1 if mono else 2 59 | format_strings = { 60 | np.float64: 'f64le', 61 | np.float32: 'f32le', 62 | np.int16: 's16le', 63 | np.int32: 's32le', 64 | np.uint32: 'u32le' 65 | } 66 | format_string = format_strings[in_type] 67 | command = [ 68 | 'ffmpeg', 69 | '-i', filename, 70 | '-f', format_string, 71 | '-acodec', 'pcm_' + format_string, 72 | '-ar', str(sr), 73 | '-ac', str(channels), 74 | '-'] 75 | p = sp.Popen(command, stdout=sp.PIPE, stderr=DEVNULL) 76 | raw, err = p.communicate() 77 | audio = np.frombuffer(raw, dtype=in_type) 78 | 79 | if channels > 1: 80 | audio = audio.reshape((-1, channels)).transpose() 81 | 82 | if audio.size == 0: 83 | return audio.astype(out_type), sr 84 | 85 | audio = convert_bit_depth(audio, in_type, out_type, normalize) 86 | 87 | return audio, sr 88 | 89 | def auwrite(fn, audio, sr): 90 | if len(audio.shape) > 1: 91 | channels = np.min(audio.shape) 92 | else: 93 | channels = 1 94 | format_strings = { 95 | 'float64': 'f64le', 96 | 'float32': 'f32le', 97 | 'int16': 's16le', 98 | 'int32': 's32le', 99 | 'uint32': 'u32le' 100 | } 101 | format_strings = {np.dtype(key): value for key,value in format_strings.items()} 102 | format_string = format_strings[audio.dtype] 103 | command = [ 104 | 'ffmpeg', 105 | '-y', 106 | '-ac', str(channels), 107 | '-ar', str(sr), 108 | '-f', format_string, 109 | '-i', 'pipe:', 110 | fn] 111 | p = sp.Popen(command, stdin=sp.PIPE, stdout=None, stderr=None) 112 | if channels > 1 and audio.shape[0] == channels: 113 | raw, err = p.communicate(audio.T.tobytes()) 114 | else: 115 | raw, err = p.communicate(audio.tobytes()) 116 | 117 | def auchannels(y): 118 | if len(y.shape) > 1: 119 | return y.shape[0] 120 | return 1 121 | 122 | def aulen(y): 123 | if len(y.shape) > 1: 124 | return y.shape[1] 125 | return len(y) 126 | 127 | def convert_fraction_to_real(framerate): 128 | if '/' in framerate: 129 | num,div = framerate.split('/') 130 | return float(num) / float(div) 131 | return float(framerate) 132 | 133 | import json 134 | def vidreadmeta(fn): 135 | if not os.path.exists(fn): 136 | raise FileNotFoundError 137 | probe = ffmpeg.probe(fn) 138 | for stream in probe['streams']: 139 | if stream['codec_type'] == 'video': 140 | meta = { 141 | 'width': int(stream['width']), 142 | 'height': int(stream['height']), 143 | 'duration': float(probe['format']['duration']), 144 | 'framerate': convert_fraction_to_real(stream['avg_frame_rate']) 145 | } 146 | return meta 147 | return None 148 | 149 | def vidread(fn, samples=None, rate=None, hwaccel=None): 150 | if not os.path.exists(fn): 151 | raise FileNotFoundError 152 | probe = ffmpeg.probe(fn) 153 | out_params = {} 154 | for stream in probe['streams']: 155 | if stream['codec_type'] == 'video': 156 | width, height = stream['width'], stream['height'] 157 | try: 158 | if stream['tags']['rotate'] in ['90','270','-90']: # not sure if -90 ever happens 159 | width, height = height, width 160 | except KeyError: 161 | pass 162 | if samples is not None: 163 | duration = float(stream['duration']) 164 | interval = duration / samples 165 | out_params['r'] = 1 / interval 166 | out_params['ss'] = interval / 2 167 | elif rate is not None: 168 | out_params['r'] = rate 169 | out_params['ss'] = 1 / (2 * rate) 170 | in_params = {} 171 | if hwaccel is not None: 172 | in_params['hwaccel'] = hwaccel 173 | channels = 3 174 | frame_number = -1 175 | try: 176 | proc = ( 177 | ffmpeg 178 | .input(fn, **in_params) 179 | .output('pipe:', format='rawvideo', pix_fmt='rgb24', **out_params) 180 | .run_async(pipe_stdout=True) 181 | ) 182 | while True: 183 | in_bytes = proc.stdout.read(width*height*channels) 184 | frame_number += 1 185 | if not in_bytes: 186 | break 187 | in_frame = ( 188 | np 189 | .frombuffer(in_bytes, np.uint8) 190 | .reshape([height, width, channels]) 191 | ) 192 | yield in_frame 193 | finally: 194 | proc.stdout.close() 195 | proc.wait() 196 | 197 | class VideoWriter: 198 | def __init__(self, fn, vcodec='libx264', fps=60, in_pix_fmt='rgb24', out_pix_fmt='yuv420p', input_args=None, output_args=None): 199 | self.fn = fn 200 | self.process = None 201 | self.input_args = {} if input_args is None else input_args 202 | self.output_args = {} if output_args is None else output_args 203 | self.input_args['framerate'] = fps 204 | self.input_args['pix_fmt'] = in_pix_fmt 205 | self.output_args['pix_fmt'] = out_pix_fmt 206 | self.output_args['vcodec'] = vcodec 207 | 208 | def add(self, frame): 209 | if self.process is None: 210 | h,w = frame.shape[:2] 211 | self.process = ( 212 | ffmpeg 213 | .input('pipe:', format='rawvideo', s='{}x{}'.format(w, h), **self.input_args) 214 | .output(self.fn, **self.output_args) 215 | .overwrite_output() 216 | .run_async(pipe_stdin=True) 217 | ) 218 | self.process.stdin.write( 219 | frame 220 | .astype(np.uint8) 221 | .tobytes() 222 | ) 223 | 224 | def close(self): 225 | if self.process is None: 226 | return 227 | self.process.stdin.close() 228 | self.process.wait() 229 | 230 | def vidwrite(fn, images, **kwargs): 231 | writer = VideoWriter(fn, **kwargs) 232 | for image in images: 233 | writer.add(image) 234 | writer.close() 235 | -------------------------------------------------------------------------------- /geometry.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import math 3 | 4 | # angle between two vectors 5 | def vector_vector_angle(v1, v2): 6 | a1 = np.arctan2(v1[1], v1[0]) 7 | a2 = np.arctan2(v2[1], v2[0]) 8 | sign = 1 if a1 > a2 else -1 9 | angle = a1 - a2 10 | K = -sign * np.pi * 2 11 | if np.abs(K + angle) < np.abs(angle): 12 | angle += K 13 | return angle 14 | 15 | # pair of points on lines p1,p2 and p3,p4 closest to each other 16 | def line_line_closest(p1,p2,p3,p4,check=False): 17 | p1 = np.asarray(p1) 18 | p2 = np.asarray(p2) 19 | p3 = np.asarray(p3) 20 | p4 = np.asarray(p4) 21 | 22 | p13 = p1 - p3 23 | p43 = p4 - p3 24 | p21 = p2 - p1 25 | 26 | d1343 = np.dot(p13, p43) 27 | d4321 = np.dot(p43, p21) 28 | d1321 = np.dot(p13, p21) 29 | d4343 = np.dot(p43, p43) 30 | d2121 = np.dot(p21, p21) 31 | 32 | denom = d2121 * d4343 - d4321 * d4321 33 | if check and abs(denom) == 0: 34 | return None 35 | numer = d1343 * d4321 - d1321 * d4343 36 | 37 | mua = numer / denom 38 | mub = (d1343 + d4321 * (mua)) / d4343 39 | 40 | if check: 41 | if mua < 0 or mua > 1 or mub < 0 or mub > 1: 42 | return None 43 | 44 | pa = p1 + mua * p21 45 | pb = p3 + mub * p43 46 | 47 | return pa, pb 48 | 49 | # point on line p1,p2 closest to point p3 50 | def line_point_closest(p1, p2, p3): 51 | p1 = np.asarray(p1) 52 | p2 = np.asarray(p2) 53 | p3 = np.asarray(p3) 54 | u = ((p3 - p1)*(p2 - p1)).sum() / np.square(np.linalg.norm(p2 - p1)) 55 | return p1 + u * (p2 - p1) 56 | 57 | # point on ray p1->p2 closest to point p3 58 | def ray_point_closest(p1, p2, p3): 59 | p1 = np.asarray(p1) 60 | p2 = np.asarray(p2) 61 | p3 = np.asarray(p3) 62 | u = ((p3 - p1)*(p2 - p1)).sum() / np.square(np.linalg.norm(p2 - p1)) 63 | u = max(u, 0) 64 | return p1 + u * (p2 - p1) 65 | 66 | # numpy multiplication order is "reversed" from openFrameworks, 67 | # or the matrices can be transposed: A*B = B^T*A 68 | 69 | # from 0/0 to width/height 70 | # to -1,+1 to +1/-1 (y axis is flipped) 71 | # note that depth is not in world units, but normalized by z_near/z_far 72 | # designed for multiple points 73 | def screen_to_world(screen, depth, viewport, extrinsics, camera_matrix): 74 | principal_point = camera_matrix[:2,2] 75 | screen -= principal_point - (viewport - 1) / 2 76 | n = len(screen) 77 | camera = np.array([ 78 | 2 * screen[:,0] / viewport[0] - 1, 79 | 1 - 2 * screen[:,1] / viewport[1], 80 | [depth] * n 81 | ]) 82 | camera = camera.T 83 | return camera_to_world(camera, extrinsics, viewport, camera_matrix) 84 | 85 | def to_homogenous(points): 86 | ones = np.ones((len(points), 1), points.dtype) 87 | return np.hstack((points, ones)) 88 | 89 | def from_homogenous(xyzw): 90 | return xyzw[:,:3] / xyzw[:,3].reshape(-1,1) 91 | 92 | def camera_to_world(camera, extrinsics, viewport, camera_matrix): 93 | mvp_matrix = get_model_view_projection_matrix(extrinsics, viewport, camera_matrix) 94 | world = np.matmul(to_homogenous(camera), np.linalg.inv(mvp_matrix)) 95 | return from_homogenous(world) 96 | 97 | def get_model_view_projection_matrix(extrinsics, viewport, camera_matrix): 98 | return np.matmul(get_model_view_matrix(extrinsics), get_projection_matrix(viewport, camera_matrix)) 99 | 100 | def get_model_view_matrix(extrinsics): 101 | return np.linalg.inv(extrinsics.T) 102 | 103 | def perspective(fovy, aspect, z_near, z_far): 104 | tan_half_fovy = math.tan(fovy / 2) 105 | result = np.zeros((4,4)) 106 | result[0,0] = 1 / (aspect * tan_half_fovy) 107 | result[1,1] = 1 / tan_half_fovy 108 | result[2,3] = -1 109 | result[2,2] = -(z_far + z_near) / (z_far - z_near) 110 | result[3,2] = -(2 * z_far * z_near) / (z_far - z_near) 111 | return result 112 | 113 | def get_projection_matrix(viewport, camera_matrix, z_near=1, z_far=1000): 114 | aspect = viewport[0] / viewport[1] 115 | fovy = 2 * math.atan(viewport[1] / (2 * camera_matrix[1,1])) 116 | return perspective(fovy, aspect, z_near, z_far) 117 | 118 | def vec_to_extrinsics(rvec,tvec): 119 | R, jacobian = cv2.Rodrigues(np.mat(rvec)) 120 | arr = np.array([[R[0,0],R[0,1],R[0,2],tvec[0]], 121 | [R[1,0],R[1,1],R[1,2],tvec[1]], 122 | [R[2,0],R[2,1],R[2,2],tvec[2]], 123 | [0,0,0,1]]) 124 | mat = np.mat(arr) 125 | extrinsics = np.transpose(np.linalg.inv(mat)) 126 | return extrinsics 127 | -------------------------------------------------------------------------------- /histogram_equalize.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def histogram_equalize(data, max_val=None, endpoint=False): 4 | input_shape = np.shape(data) 5 | data_flat = np.asarray(data).flatten() 6 | if max_val is None: 7 | max_val = data_flat.max() 8 | indices = np.argsort(data_flat) 9 | replacements = np.linspace(0, max_val, len(indices), endpoint=endpoint) 10 | data_flat[indices] = replacements 11 | return data_flat.reshape(*input_shape) -------------------------------------------------------------------------------- /imutil.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import numpy as np 4 | import PIL.Image 5 | import shutil 6 | from utils.color_conversion import to_single_rgb, to_single_gray, rb_swap 7 | 8 | try: # Python 2 9 | from cStringIO import StringIO as BytesIO 10 | except: # Python 3 11 | from io import BytesIO 12 | 13 | # should add code to automatically scale 0-1 to 0-255 14 | def imshow(img, fmt='png', retina=False, zoom=None): 15 | import IPython.display 16 | if img is None: 17 | raise TypeError('input image not provided') 18 | 19 | if isinstance(img, str): 20 | if img.startswith('http:') or img.startswith('https:'): 21 | IPython.display.display(IPython.display.Image(url=img, retina=retina)) 22 | else: 23 | IPython.display.display(IPython.display.Image(filename=img, retina=retina)) 24 | return 25 | 26 | if len(img.shape) == 1: 27 | n = len(img) 28 | side = int(np.sqrt(n)) 29 | if (side * side) == n: 30 | img = img.reshape(side, side) 31 | else: 32 | raise ValueError('input is one-dimensional', img.shape) 33 | if len(img.shape) == 3 and img.shape[-1] == 1: 34 | img = img.squeeze() 35 | img = np.uint8(np.clip(img, 0, 255)) 36 | if fmt == 'jpg': 37 | fmt = 'jpeg' 38 | if fmt == 'jpeg': 39 | img = to_single_rgb(img) 40 | image_data = BytesIO() 41 | PIL.Image.fromarray(img).save(image_data, fmt) 42 | height, width = img.shape[:2] 43 | if zoom is not None: 44 | width *= zoom 45 | height *= zoom 46 | IPython.display.display(IPython.display.Image(data=image_data.getvalue(), 47 | width=width, 48 | height=height, 49 | retina=retina)) 50 | 51 | # jpeg4py is 2x as fast as opencv for jpegs, but more unstable 52 | def imread(filename, mode=None, ext=None): 53 | if ext is None: 54 | _, ext = os.path.splitext(filename) 55 | ext = ext.lower() 56 | img = cv2.imread(filename, cv2.IMREAD_UNCHANGED) 57 | if img is not None: 58 | if len(img.shape) > 2: 59 | img = rb_swap(img) 60 | if img is not None: 61 | if mode == 'rgb': 62 | img = to_single_rgb(img) 63 | elif mode == 'gray': 64 | img = to_single_gray(img) 65 | return img 66 | 67 | def imwrite(filename, img): 68 | if img is None: 69 | return 70 | if len(img.shape) == 2: 71 | return cv2.imwrite(filename, img) 72 | if len(img.shape) == 3: 73 | if img.shape[-1] == 1: 74 | return cv2.imwrite(filename, img[:,:,0]) 75 | elif img.shape[-1] == 3: 76 | return cv2.imwrite(filename, img[...,::-1]) 77 | elif img.shape[-1] == 4: 78 | return cv2.imwrite(filename, img[...,(2,1,0,3)]) 79 | else: 80 | raise Exception('Unsupported number of channels for shape', img.shape) 81 | else: 82 | raise Exception('Unsupported image shape', img.shape) 83 | 84 | def downsample(img, scale=None, output_wh=None, max_side=None, min_side=None, block_size=None, mode=None): 85 | if max_side is not None: 86 | cur_max_side = max(img.shape[:2]) 87 | scale = max_side / cur_max_side 88 | if min_side is not None: 89 | cur_min_side = min(img.shape[:2]) 90 | scale = min_side / cur_min_side 91 | if scale is not None: 92 | output_wh = (int(np.round(img.shape[1]*scale)), 93 | int(np.round(img.shape[0]*scale))) 94 | if block_size is not None: 95 | output_wh = (img.shape[1]//block_size, img.shape[0]//block_size) 96 | else: 97 | block_size = img.shape[1]//output_wh[0] 98 | if block_size > 1: 99 | img = cv2.blur(img, (block_size, block_size)) 100 | return cv2.resize(img, output_wh, interpolation=cv2.INTER_AREA if mode is None else mode) 101 | 102 | def upsample(img, scale=None, output_wh=None, max_side=None, min_side=None, mode=None): 103 | if max_side is not None: 104 | cur_max_side = max(img.shape[:2]) 105 | scale = max_side / cur_max_side 106 | if min_side is not None: 107 | cur_min_side = min(img.shape[:2]) 108 | scale = min_side / cur_min_side 109 | if output_wh is None: 110 | output_wh = (int(np.round(img.shape[1]*scale)), 111 | int(np.round(img.shape[0]*scale))) 112 | return cv2.resize(img, output_wh, interpolation=cv2.INTER_CUBIC if mode is None else mode) 113 | 114 | # output_wh value None in one dimension means "scale proportionally to other dimension" 115 | # output_wh value -1 in one dimension means "use the existing value" 116 | def imresize(img, scale=None, output_wh=None, max_side=None, min_side=None, mode=None): 117 | big = True 118 | if max_side is not None: 119 | cur_max_side = max(img.shape[:2]) 120 | big = max_side > cur_max_side 121 | elif min_side is not None: 122 | cur_min_side = min(img.shape[:2]) 123 | big = min_side > cur_min_side 124 | elif output_wh is not None: 125 | if output_wh[0] is None: 126 | scale = output_wh[1] / img.shape[0] 127 | output_wh = None 128 | elif output_wh[1] is None: 129 | scale = output_wh[0] / img.shape[1] 130 | output_wh = None 131 | elif output_wh[0] == -1: 132 | output_wh[0] = img.shape[1] 133 | elif output_wh[1] == -1: 134 | output_wh[1] = img.shape[0] 135 | if output_wh is not None: 136 | big = output_wh[0] > img.shape[1] 137 | if scale is not None: 138 | big = scale > 1 139 | 140 | if big: 141 | return upsample(img, scale, output_wh, max_side, min_side, mode) 142 | else: 143 | return downsample(img, scale, output_wh, max_side, min_side, mode) 144 | -------------------------------------------------------------------------------- /imwarp.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | 4 | def translate_image(img, xy, borderValue=(128,128,128)): 5 | h,w = img.shape[:2] 6 | if len(img.shape) == 2 and hasattr(borderValue, '__len__'): 7 | borderValue = borderValue[0] 8 | x,y = xy 9 | matrix = np.float32([[1,0,x],[0,1,y]]) 10 | return cv2.warpAffine(img, matrix, (w,h), 11 | flags=cv2.INTER_CUBIC, 12 | borderMode=cv2.BORDER_CONSTANT, 13 | borderValue=borderValue) 14 | 15 | # pivot is y,x pixels, pivot_percent in ratio to height and width 16 | def rotate_image(img, angle, pivot_percent=(0.5, 0.5), pivot=None, scale=1.0, borderValue=(128,128,128)): 17 | h,w = img.shape[:2] 18 | if len(img.shape) == 2 and hasattr(borderValue, '__len__'): 19 | borderValue = borderValue[0] 20 | if pivot is None: 21 | pivot = (int(h * pivot_percent[0]), int(w * pivot_percent[1])) 22 | matrix = cv2.getRotationMatrix2D((pivot[1], pivot[0]), angle, scale) # uses x, y pivot 23 | return cv2.warpAffine(img, matrix, (w, h), 24 | flags=cv2.INTER_CUBIC, 25 | borderMode=cv2.BORDER_CONSTANT, 26 | borderValue=borderValue) -------------------------------------------------------------------------------- /itertools.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from itertools import islice 3 | from itertools import chain 4 | 5 | def np_chunks(x, chunk_size): 6 | chunk_count = len(x)//chunk_size 7 | n = chunk_count * chunk_size 8 | shape = x.shape[1:] 9 | return x[:n].reshape(chunk_count, chunk_size, *shape) 10 | 11 | def chunks(x, n): 12 | # return slices of lists 13 | if hasattr(x, '__len__'): 14 | for i in range(0, len(x), n): 15 | yield x[i:i+n] 16 | else: 17 | # return sub-generators of generators 18 | i = iter(x) 19 | for e in i: 20 | yield chain([e], islice(i, n-1)) -------------------------------------------------------------------------------- /license.md: -------------------------------------------------------------------------------- 1 | Copyright (c) 2018 Kyle McDonald 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining 4 | a copy of this software and associated documentation files (the 5 | "Software"), to deal in the Software without restriction, including 6 | without limitation the rights to use, copy, modify, merge, publish, 7 | distribute, sublicense, and/or sell copies of the Software, and to 8 | permit persons to whom the Software is furnished to do so, subject to 9 | the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be 12 | included in all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /list_files.py: -------------------------------------------------------------------------------- 1 | import os 2 | import fnmatch 3 | 4 | # does not operate recursively 5 | def list_directories(directory, exclude_prefixes=('.',)): 6 | for f in os.listdir(directory): 7 | if f.startswith(exclude_prefixes): 8 | continue 9 | joined = os.path.join(directory, f) 10 | if os.path.isdir(joined): 11 | yield joined 12 | 13 | # extensions can be a string or list, can include the preceding . or not 14 | # operates recursively 15 | def list_files(directory, extensions=None, exclude_prefixes=('.',)): 16 | if type(extensions) == str: 17 | extensions = [extensions] 18 | if extensions is not None: 19 | extensions = [('' if e.startswith('.') else '.') + e for e in extensions] 20 | for root, dirnames, filenames in os.walk(directory): 21 | filenames = [f for f in filenames if not f.startswith(exclude_prefixes)] 22 | dirnames[:] = [d for d in dirnames if not d.startswith(exclude_prefixes)] 23 | for filename in filenames: 24 | base, ext = os.path.splitext(filename) 25 | joined = os.path.join(root, filename) 26 | if extensions is None or ext.lower() in extensions: 27 | yield joined 28 | 29 | def get_stem(fn): 30 | return os.path.splitext(os.path.basename(fn))[0] -------------------------------------------------------------------------------- /memoize.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import pickle 3 | import inspect 4 | import os 5 | 6 | def memoize(func): 7 | is_method = 'self' in inspect.getfullargspec(func).args 8 | def wrapper(*args, **kwargs): 9 | args_to_hash = args[1:] if is_method else args 10 | h = hashlib.md5(str(args_to_hash).encode()).hexdigest() 11 | fn = f"cache/{func.__name__}_{h}.pkl" 12 | try: 13 | os.makedirs('cache', exist_ok=True) 14 | with open(fn, 'rb') as f: 15 | return pickle.load(f) 16 | except FileNotFoundError: 17 | result = func(*args, **kwargs) 18 | with open(fn, 'wb') as f: 19 | pickle.dump(result, f) 20 | return result 21 | return wrapper -------------------------------------------------------------------------------- /mosaic.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import math 3 | 4 | def find_rectangle(n): 5 | max_side = int(math.sqrt(n)) 6 | for h in range(2, max_side+1)[::-1]: 7 | w = n // h 8 | if (h * w) == n: 9 | return (h, w) 10 | return (n, 1) 11 | 12 | def swapaxes(x,a,b): 13 | try: 14 | return x.swapaxes(a,b) 15 | except AttributeError: # support pytorch 16 | return x.transpose(a,b) 17 | 18 | # 1d images (n, h*w): no 19 | # 2d images (n, h, w): yes 20 | # 3d images (n, h, w, c): yes 21 | def make_mosaic(x, nx=None, ny=None): 22 | if not isinstance(x, np.ndarray): 23 | x = np.asarray(x) 24 | 25 | n, h, w = x.shape[:3] 26 | has_channels = len(x.shape) > 3 27 | if has_channels: 28 | c = x.shape[3] 29 | 30 | if nx is None and ny is None: 31 | ny,nx = find_rectangle(n) 32 | elif ny is None: 33 | ny = n//nx 34 | elif nx is None: 35 | nx = n//ny 36 | 37 | end_shape = (w,c) if has_channels else (w,) 38 | mosaic = x.reshape(ny, nx, h, *end_shape) 39 | mosaic = swapaxes(mosaic, 1, 2) 40 | hh = mosaic.shape[0] * mosaic.shape[1] 41 | ww = mosaic.shape[2] * mosaic.shape[3] 42 | end_shape = (ww,c) if has_channels else (ww,) 43 | mosaic = mosaic.reshape(hh, *end_shape) 44 | return mosaic 45 | 46 | # 1d images (n, h*w): no 47 | # 2d images (n, h, w): yes 48 | # 3d images (n, h, w, c): yes 49 | # assumes images are square if underspecified 50 | def unmake_mosaic(mosaic, nx=None, ny=None, w=None, h=None): 51 | hh, ww = mosaic.shape[:2] 52 | 53 | if nx is not None or ny is not None: 54 | if nx is None: 55 | h = hh//ny 56 | w = h 57 | nx = ww//w 58 | elif ny is None: 59 | w = ww//nx 60 | h = w 61 | ny = hh//h 62 | else: 63 | w = ww//nx 64 | h = hh//ny 65 | 66 | elif w is not None or h is not None: 67 | if w is None: 68 | w = h 69 | elif h is None: 70 | h = w 71 | nx = ww//w 72 | ny = hh//h 73 | 74 | end_shape = (w, mosaic.shape[2]) if len(mosaic.shape) > 2 else (w,) 75 | 76 | x = mosaic.reshape(ny, h, nx, *end_shape) 77 | x = swapaxes(x, 1, 2) 78 | x = x.reshape(-1, h, *end_shape) 79 | return x -------------------------------------------------------------------------------- /multiworker.py: -------------------------------------------------------------------------------- 1 | from threading import Thread 2 | from multiprocessing import Process 3 | from multiprocessing import Queue 4 | 5 | class MultiWorker: 6 | def __init__(self, worker_type, job=None, job_loop=None, num_workers=1): 7 | self.input_queue = Queue() 8 | self.output_queue = Queue() 9 | self.workers = [] 10 | if job_loop is None: 11 | def job_loop(id, input, output, ready): 12 | ready.put(id) 13 | while True: 14 | task = input.get() 15 | if task is None: 16 | break 17 | result = job(task) 18 | output.put(result) 19 | ready = Queue() 20 | for i in range(num_workers): 21 | worker = Process(target=job_loop, args=(i, self.input_queue, self.output_queue, ready)) 22 | worker.start() 23 | self.workers.append(worker) 24 | # block until all workers are ready 25 | for i in range(num_workers): 26 | ready.get() 27 | 28 | def put(self, task): 29 | self.input_queue.put(task) 30 | 31 | def get(self): 32 | return self.output_queue.get() 33 | 34 | def join(self): 35 | for worker in self.workers: 36 | self.input_queue.put(None) 37 | for worker in self.workers: 38 | worker.join() 39 | 40 | class MultiThreadWorker(MultiWorker): 41 | def __init__(self, **kwargs): 42 | super().__init__(Thread, **kwargs) 43 | 44 | class MultiProcessWorker(MultiWorker): 45 | def __init__(self, **kwargs): 46 | super().__init__(Process, **kwargs) -------------------------------------------------------------------------------- /pixelated.py: -------------------------------------------------------------------------------- 1 | from IPython.core.display import display_html, HTML 2 | 3 | def pixelate(): 4 | display_html(HTML('')) -------------------------------------------------------------------------------- /plot_images.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def plot_images(images, xy, blend=np.maximum, canvas_shape=(512,512), fill=0): 4 | h,w = images.shape[1:3] 5 | if images.ndim == 4: 6 | canvas_shape = (canvas_shape[0], canvas_shape[1], images.shape[3]) 7 | 8 | min_xy = np.amin(xy, 0) 9 | max_xy = np.amax(xy, 0) 10 | 11 | min_canvas = np.array((0, 0)) 12 | max_canvas = np.array((canvas_shape[0] - h, canvas_shape[1] - w)) 13 | 14 | xy_mapped = min_canvas + (xy - min_xy) * (max_canvas - min_canvas) / (max_xy - min_xy) 15 | xy_mapped = xy_mapped.astype(int) 16 | 17 | canvas = np.full(canvas_shape, fill) 18 | for image, pos in zip(images, xy_mapped): 19 | x_off, y_off = pos 20 | sub_canvas = canvas[y_off:y_off+h, x_off:x_off+w] 21 | sub_image = image[:h, :w] 22 | try: 23 | canvas[y_off:y_off+h, x_off:x_off+w] = blend(sub_canvas, sub_image) 24 | except ValueError: 25 | print(pos, h, w, min_canvas, max_canvas) 26 | raise 27 | 28 | return canvas -------------------------------------------------------------------------------- /progress.py: -------------------------------------------------------------------------------- 1 | from multiprocessing import Pool, cpu_count 2 | import time 3 | from datetime import datetime, timedelta 4 | import sys 5 | 6 | try: 7 | from IPython.display import clear_output 8 | except ModuleNotFoundError: 9 | def clear_output(*args, **kwargs): 10 | pass 11 | 12 | def progress(itr, total=None, update_interval=1, clear=True): 13 | if total is None and hasattr(itr, '__len__'): 14 | total = len(itr) 15 | if total == 0: 16 | return 17 | if total: 18 | print(f'0/{total} 0s 0/s') 19 | else: 20 | print('0 0s 0/s') 21 | start_time = None 22 | last_time = None 23 | for i, x in enumerate(itr): 24 | cur_time = time.time() 25 | if start_time is None: 26 | start_time = cur_time 27 | last_time = cur_time 28 | yield x 29 | if cur_time - last_time > update_interval: 30 | duration = cur_time - start_time 31 | speed = (i + 1) / duration 32 | duration_str = timedelta(seconds=round(duration)) 33 | if clear: 34 | clear_output(wait=True) 35 | if total: 36 | duration_total = duration * total / (i + 1) 37 | duration_remaining = duration_total - duration 38 | duration_remaining_str = timedelta(seconds=round(duration_remaining)) 39 | pct = 100. * (i + 1) / total 40 | print(f'{pct:.2f}% {i+1}/{total} {duration_str}<{duration_remaining_str} {speed:.2f}/s') 41 | else: 42 | print(f'{i+1} {duration_str} {speed:.2f}/s') 43 | last_time = cur_time 44 | 45 | duration = time.time() - start_time 46 | speed = (i + 1) / duration 47 | duration_str = timedelta(seconds=round(duration)) 48 | if clear: 49 | clear_output(wait=True) 50 | print(f'{i+1} {duration_str} {speed:.2f}/s') 51 | 52 | class job_wrapper(object): 53 | def __init__(self, job): 54 | self.job = job 55 | def __call__(self, args): 56 | i, task = args 57 | return i, self.job(task) 58 | 59 | def progress_parallel(job, tasks, total=None, processes=None, **kwargs): 60 | if processes == 1: 61 | return [job(task) for task in progress(tasks)] 62 | 63 | results = [] 64 | if total is None and hasattr(tasks, '__len__'): 65 | total = len(tasks) 66 | if processes is None: 67 | processes = cpu_count() 68 | try: 69 | with Pool(processes) as pool: 70 | results = list(progress(pool.imap_unordered(job_wrapper(job), enumerate(tasks)), 71 | total=total, **kwargs)) 72 | results.sort() 73 | return [x for i,x in results] 74 | except KeyboardInterrupt: 75 | pass 76 | 77 | from joblib import Parallel, delayed 78 | def progress_parallel_joblib(job, tasks, total=None, processes=None, backend=None): 79 | if total is None and hasattr(tasks, '__len__'): 80 | total = len(tasks) 81 | if processes is None: 82 | processes = -1 83 | 84 | try: 85 | results = Parallel(n_jobs=processes)(delayed(job)(task) for task in progress(tasks)) 86 | return results 87 | except KeyboardInterrupt: 88 | pass -------------------------------------------------------------------------------- /rainbow.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def build_rainbow(n, curve=None): 4 | rgb = [] 5 | width = 2 * np.pi 6 | for i in range(3): 7 | offset = -i * width / 3 8 | cur = np.cos(np.linspace(offset, offset + width, n)) 9 | rgb.append(cur) 10 | rainbow = (1 + np.vstack(rgb)) / 2 11 | if curve: 12 | rainbow = curve(rainbow) 13 | rainbow = np.minimum(rainbow * 256, 255).astype(np.uint8) 14 | return rainbow.T 15 | 16 | def to_rainbow(x, n=1024, minimum=None, maximum=None): 17 | rainbow_lookup = build_rainbow(n) 18 | indices = np.copy(x).astype('float') 19 | if minimum is None: 20 | minimum = x.min() 21 | if maximum is None: 22 | maximum = x.max() 23 | indices -= minimum 24 | indices *= n / (maximum - minimum) 25 | indices = np.minimum(indices, n - 1).astype(np.int) 26 | return rainbow_lookup.take(indices, axis=0) 27 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # python-utils 2 | 3 | Disorganized collection of useful functions for working with audio and images, especially in the context of machine learning. 4 | 5 | For Python 3.7+ 6 | 7 | Install using pip: 8 | 9 | ``` 10 | pip install git+git://github.com/kylemcdonald/python-utils.git 11 | ``` 12 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!usr/bin/env python 2 | 3 | from setuptools import setup, find_packages 4 | 5 | setup(name="pyutils", 6 | version="0.1", 7 | description="Test for pip install git+", 8 | url="https://github.com/kylemcdonald/python-utils", 9 | install_requires=["numpy", "opencv-python", "Pillow"], 10 | packages=find_packages() 11 | ) 12 | -------------------------------------------------------------------------------- /specshow.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | from utils.imutil import imshow 4 | 5 | # expects output directly from librosa.stft or librosa.cqt 6 | def specshow(x, sr=44100, hop_length=512, max_frames=1960, skip=1, gamma=6, use_mag=True, cmap='inferno', zoom=None, show=True): 7 | bins, frames = x.shape 8 | seconds = (frames * hop_length) / sr 9 | minutes = int(seconds // 60) 10 | seconds = seconds - (minutes * 60) 11 | sx = x[:,:max_frames*skip:skip] 12 | if not np.iscomplexobj(sx): 13 | spec = np.copy(sx) 14 | else: 15 | mag = np.abs(sx) 16 | spec = mag if use_mag else np.log(1+mag**2) 17 | spec -= spec.min() 18 | spec = spec / spec.max() 19 | spec **= 1 / gamma 20 | cm = plt.get_cmap(cmap) 21 | spec = cm(spec)[:,:,:3] 22 | img = 255 * np.flipud(spec) 23 | if show: 24 | imshow(img, retina=zoom is None, zoom=zoom) 25 | print(f'{minutes}:{seconds:04.2f} @ {sr}Hz, {frames} frames x {bins} bins @ {hop_length} hop_length') 26 | else: 27 | return img -------------------------------------------------------------------------------- /stratified.py: -------------------------------------------------------------------------------- 1 | import collections 2 | import itertools 3 | 4 | def dict_to_iterators(x): 5 | if isinstance(x, dict): 6 | return itertools.cycle([dict_to_iterators(x[e]) for e in sorted(x.keys())]) 7 | else: 8 | return itertools.cycle(x) 9 | 10 | def deep_default_dict(n): 11 | if n == 1: 12 | return collections.defaultdict(list) 13 | else: 14 | return collections.defaultdict(lambda: deep_default_dict(n-1)) 15 | 16 | def append_deep_default_dict(d, keys, value): 17 | cur = d 18 | for key in keys: 19 | cur = cur[key] 20 | cur.append(value) 21 | 22 | def build_stratified(values, selected=None): 23 | x = deep_default_dict(len(values)) 24 | for i, keys in enumerate(zip(*values)): 25 | if selected is None or i in selected: 26 | append_deep_default_dict(x, keys, i) 27 | return dict_to_iterators(x) 28 | 29 | def next_recursive(x): 30 | if isinstance(x, collections.Iterable): 31 | return next_recursive(next(x)) 32 | else: 33 | return x -------------------------------------------------------------------------------- /timing.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | from time import sleep 3 | import contextlib 4 | import numpy as np 5 | from IPython.display import clear_output 6 | import json 7 | from multiprocessing import Queue 8 | from multiprocessing import Process 9 | from threading import get_ident 10 | import os 11 | import sys 12 | 13 | class Ticker(): 14 | def __init__(self, update_rate=1): 15 | self.total_ticks = 0 16 | self.start_time = None 17 | self.last_print = None 18 | self.recent = [] 19 | self.clear_output = 'ipykernel' in sys.modules 20 | self.update_rate = update_rate 21 | 22 | def tick(self): 23 | cur_time = time() 24 | self.recent.append(cur_time) 25 | if self.start_time is not None: 26 | duration = cur_time - self.start_time 27 | cur_print = int(duration) 28 | if cur_print != self.last_print and cur_print % self.update_rate == 0: 29 | fps = self.total_ticks / duration 30 | denominator = (self.recent[-1] - self.recent[0]) 31 | if denominator > 0: 32 | recent_fps = (len(self.recent) - 1) / denominator 33 | else: 34 | recent_fps = 0 35 | jitter = np.std(self.recent) 36 | print(f'recent: {recent_fps:0.2f} fps, all: {fps:0.2f} fps, jitter:', format_time(jitter)) 37 | self.last_print = cur_print 38 | self.recent = [] 39 | if self.clear_output: 40 | clear_output(wait=True) 41 | else: 42 | self.start_time = cur_time 43 | self.total_ticks += 1 44 | 45 | class Profiler(): 46 | def __init__(self, name): 47 | self.timing = {} 48 | self.last_print = None 49 | self.name = name 50 | 51 | @contextlib.contextmanager 52 | def profile(self, name): 53 | start = time() 54 | yield 55 | elapsed = time() - start 56 | if name in self.timing: 57 | self.timing[name][0] += elapsed 58 | self.timing[name][1] += 1 59 | else: 60 | self.timing[name] = [elapsed, 1] 61 | 62 | def print(self): 63 | cur = int(time()) 64 | if cur != self.last_print: 65 | parts = [] 66 | for name in sorted(self.timing): 67 | elapsed, count = self.timing[name] 68 | average = format_time(elapsed / count) 69 | parts.append(f'{name}: {average}') 70 | print(self.name + ': ' + ', '.join(parts)) 71 | self.last_print = cur 72 | self.timing = {} 73 | return True 74 | return False 75 | 76 | def format_time(seconds): 77 | if seconds < 1/1e5: 78 | return f'{seconds*1e6:0.1f}us' 79 | if seconds < 1/1e4: 80 | return f'{seconds*1e6:0.0f}us' 81 | if seconds < 1/1e2: 82 | return f'{seconds*1e3:0.1f}ms' 83 | if seconds < 1/1e1: 84 | return f'{seconds*1e3:0.0f}ms' 85 | else: 86 | return f'{seconds:0.1f}s' 87 | 88 | class Tracer(): 89 | def __init__(self, fn): 90 | self.queue = Queue() 91 | self.pid_names = {} 92 | self.tid_names = {} 93 | 94 | self.enabled = True 95 | if fn is None: 96 | self.enabled = False 97 | return 98 | 99 | def tracer_loop(): 100 | output = open(fn, 'w') 101 | output.write('[\n') 102 | while True: 103 | entry = self.queue.get() 104 | line = json.dumps(entry, separators=(',', ':')) 105 | output.write(line + ',\n') 106 | 107 | self.set_pid_name('Main') 108 | self.tracer = Process(target=tracer_loop) 109 | self.tracer.start() 110 | 111 | def set_tid_name(self, name): 112 | self.tid_names[get_ident()] = name 113 | 114 | def set_pid_name(self, name): 115 | self.pid_names[os.getpid()] = name 116 | 117 | @contextlib.contextmanager 118 | def trace(self, name, pid=None, tid=None): 119 | ts = time() 120 | yield 121 | dur = time() - ts 122 | if not self.enabled: 123 | return 124 | ts *= 1000000 # seconds to microseconds 125 | dur *= 1000000 # seconds to microseconds 126 | if pid is None: 127 | pid = os.getpid() 128 | if pid in self.pid_names: 129 | pid = self.pid_names[pid] 130 | if tid is None: 131 | tid = get_ident() 132 | if tid in self.tid_names: 133 | tid = self.tid_names[tid] 134 | entry = { 135 | 'name': name, 136 | 'ph': 'X', 137 | 'ts': ts, 138 | 'dur': dur, 139 | 'pid': pid, 140 | 'tid': tid 141 | } 142 | self.queue.put(entry) -------------------------------------------------------------------------------- /wavshow.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | 4 | def wavshow(audio, sr=None, resolution=2048, plot=True): 5 | if len(audio.shape) == 1: 6 | audio = audio.reshape(-1,1) 7 | if audio.shape[1] > audio.shape[0]: 8 | audio = audio.T 9 | duration = len(audio) 10 | if sr is not None: 11 | duration /= sr 12 | plt.figure(figsize=(16,4)) 13 | n = len(audio) 14 | k = n // resolution 15 | if k > 1: 16 | audio = audio[:(n//k)*k] 17 | audio = audio.reshape(-1,k) 18 | mean = audio.mean(axis=1) 19 | std = audio.std(axis=1) 20 | x_ticks = np.linspace(0, duration, audio.shape[0]) 21 | plt.fill_between(x_ticks, mean - std, mean + std, color='black', lw=0) 22 | plt.fill_between(x_ticks, audio.min(axis=1), audio.max(axis=1), alpha=0.5, color='gray', lw=0) 23 | else: 24 | x_ticks = np.linspace(0, duration, len(audio)) 25 | plt.plot(x_ticks, audio, color='black', lw=1) 26 | plt.xlim(0, duration) 27 | abs_max = np.abs(audio).max() 28 | plt.ylim(-abs_max, abs_max) 29 | plt.xlabel('Seconds' if sr is not None else 'Samples') 30 | if plot: 31 | plt.show() --------------------------------------------------------------------------------