├── README.md ├── common.py ├── main.py ├── tests └── test_common.py ├── tst_scene_render.py └── video.py /README.md: -------------------------------------------------------------------------------- 1 | ## Vehicle Speed Estimation 2 | 3 | ### Overview 4 | This project uses the optical flow algorithm, specifically the Lucas-Kanade tracker, to estimate vehicle speeds from mono camera (CCTV) footage. 5 | 6 | ### Prerequisites 7 | - Python 3.x 8 | - Required libraries: `opencv-python`, `numpy` 9 | 10 | ### Setup 11 | 12 | 1. **Clone the Repository** 13 | ```bash 14 | git clone https://github.com/swhan0329/vehicle_speed_estimation.git 15 | cd vehicle_speed_estimation 16 | ``` 17 | 18 | 2. **Install Dependencies** 19 | ```bash 20 | pip install -r requirements.txt 21 | ``` 22 | 23 | ### Usage 24 | 25 | 1. **With an Input Video** 26 | ```bash 27 | python main.py [input video name] 28 | ``` 29 | Replace `[input video name]` with the path to your video file. 30 | 31 | 2. **Without an Input Video** 32 | The script will automatically use the webcam on your computer. 33 | ```bash 34 | python main.py 35 | ``` 36 | 37 | ### File Descriptions 38 | 39 | - **main.py**: The main script to run the vehicle speed estimation. 40 | - **video.py**: Contains functions to handle video input and processing. 41 | - **common.py**: Includes common functions and utilities used across the project. 42 | - **tst_scene_render.py**: Test and render scenes for visualization. 43 | 44 | ### Example 45 | To run the speed estimation on a sample video: 46 | ```bash 47 | python main.py sample_video.mp4 48 | ``` 49 | 50 | To use the webcam for live speed estimation: 51 | ```bash 52 | python main.py 53 | ``` 54 | 55 | ### Additional Notes 56 | - Ensure that your video has a clear view of the road and vehicles for accurate speed estimation. 57 | - Adjust parameters in `common.py` if needed to fit specific requirements or to improve performance. 58 | 59 | ### Running Tests 60 | Execute the unit tests with Python's built-in test runner: 61 | ```bash 62 | python -m unittest discover -s tests 63 | ``` 64 | -------------------------------------------------------------------------------- /common.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | ''' 4 | This module contains some common routines used by other samples. 5 | ''' 6 | 7 | # Python 2/3 compatibility 8 | from __future__ import print_function 9 | import sys 10 | PY3 = sys.version_info[0] == 3 11 | 12 | if PY3: 13 | from functools import reduce 14 | 15 | import numpy as np 16 | import cv2 as cv 17 | 18 | # built-in modules 19 | import os 20 | import itertools as it 21 | from contextlib import contextmanager 22 | 23 | image_extensions = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.pbm', '.pgm', '.ppm'] 24 | 25 | class Bunch(object): 26 | def __init__(self, **kw): 27 | self.__dict__.update(kw) 28 | def __str__(self): 29 | return str(self.__dict__) 30 | 31 | def splitfn(fn): 32 | path, fn = os.path.split(fn) 33 | name, ext = os.path.splitext(fn) 34 | return path, name, ext 35 | 36 | def anorm2(a): 37 | return (a*a).sum(-1) 38 | def anorm(a): 39 | return np.sqrt( anorm2(a) ) 40 | 41 | def homotrans(H, x, y): 42 | xs = H[0, 0]*x + H[0, 1]*y + H[0, 2] 43 | ys = H[1, 0]*x + H[1, 1]*y + H[1, 2] 44 | s = H[2, 0]*x + H[2, 1]*y + H[2, 2] 45 | return xs/s, ys/s 46 | 47 | def to_rect(a): 48 | a = np.ravel(a) 49 | if len(a) == 2: 50 | a = (0, 0, a[0], a[1]) 51 | return np.array(a, np.float64).reshape(2, 2) 52 | 53 | def rect2rect_mtx(src, dst): 54 | src, dst = to_rect(src), to_rect(dst) 55 | cx, cy = (dst[1] - dst[0]) / (src[1] - src[0]) 56 | tx, ty = dst[0] - src[0] * (cx, cy) 57 | M = np.float64([[ cx, 0, tx], 58 | [ 0, cy, ty], 59 | [ 0, 0, 1]]) 60 | return M 61 | 62 | 63 | def lookat(eye, target, up = (0, 0, 1)): 64 | fwd = np.asarray(target, np.float64) - eye 65 | fwd /= anorm(fwd) 66 | right = np.cross(fwd, up) 67 | right /= anorm(right) 68 | down = np.cross(fwd, right) 69 | R = np.float64([right, down, fwd]) 70 | tvec = -np.dot(R, eye) 71 | return R, tvec 72 | 73 | def mtx2rvec(R): 74 | w, u, vt = cv.SVDecomp(R - np.eye(3)) 75 | p = vt[0] + u[:,0]*w[0] # same as np.dot(R, vt[0]) 76 | c = np.dot(vt[0], p) 77 | s = np.dot(vt[1], p) 78 | axis = np.cross(vt[0], vt[1]) 79 | return axis * np.arctan2(s, c) 80 | 81 | def draw_str(dst, target, s): 82 | x, y = target 83 | cv.putText(dst, s, (x, y), cv.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255),3, cv.LINE_AA) 84 | 85 | class Sketcher: 86 | def __init__(self, windowname, dests, colors_func): 87 | self.prev_pt = None 88 | self.windowname = windowname 89 | self.dests = dests 90 | self.colors_func = colors_func 91 | self.dirty = False 92 | self.show() 93 | cv.setMouseCallback(self.windowname, self.on_mouse) 94 | 95 | def show(self): 96 | cv.imshow(self.windowname, self.dests[0]) 97 | 98 | def on_mouse(self, event, x, y, flags, param): 99 | pt = (x, y) 100 | if event == cv.EVENT_LBUTTONDOWN: 101 | self.prev_pt = pt 102 | elif event == cv.EVENT_LBUTTONUP: 103 | self.prev_pt = None 104 | 105 | if self.prev_pt and flags & cv.EVENT_FLAG_LBUTTON: 106 | for dst, color in zip(self.dests, self.colors_func()): 107 | cv.line(dst, self.prev_pt, pt, color, 5) 108 | self.dirty = True 109 | self.prev_pt = pt 110 | self.show() 111 | 112 | 113 | # palette data from matplotlib/_cm.py 114 | _jet_data = {'red': ((0., 0, 0), (0.35, 0, 0), (0.66, 1, 1), (0.89,1, 1), 115 | (1, 0.5, 0.5)), 116 | 'green': ((0., 0, 0), (0.125,0, 0), (0.375,1, 1), (0.64,1, 1), 117 | (0.91,0,0), (1, 0, 0)), 118 | 'blue': ((0., 0.5, 0.5), (0.11, 1, 1), (0.34, 1, 1), (0.65,0, 0), 119 | (1, 0, 0))} 120 | 121 | cmap_data = { 'jet' : _jet_data } 122 | 123 | def make_cmap(name, n=256): 124 | data = cmap_data[name] 125 | xs = np.linspace(0.0, 1.0, n) 126 | channels = [] 127 | eps = 1e-6 128 | for ch_name in ['blue', 'green', 'red']: 129 | ch_data = data[ch_name] 130 | xp, yp = [], [] 131 | for x, y1, y2 in ch_data: 132 | xp += [x, x+eps] 133 | yp += [y1, y2] 134 | ch = np.interp(xs, xp, yp) 135 | channels.append(ch) 136 | return np.uint8(np.array(channels).T*255) 137 | 138 | def nothing(*arg, **kw): 139 | pass 140 | 141 | def clock(): 142 | return cv.getTickCount() / cv.getTickFrequency() 143 | 144 | @contextmanager 145 | def Timer(msg): 146 | print(msg, '...',) 147 | start = clock() 148 | try: 149 | yield 150 | finally: 151 | print("%.2f ms" % ((clock()-start)*1000)) 152 | 153 | class StatValue: 154 | def __init__(self, smooth_coef = 0.5): 155 | self.value = None 156 | self.smooth_coef = smooth_coef 157 | def update(self, v): 158 | if self.value is None: 159 | self.value = v 160 | else: 161 | c = self.smooth_coef 162 | self.value = c * self.value + (1.0-c) * v 163 | 164 | class RectSelector: 165 | def __init__(self, win, callback): 166 | self.win = win 167 | self.callback = callback 168 | cv.setMouseCallback(win, self.onmouse) 169 | self.drag_start = None 170 | self.drag_rect = None 171 | def onmouse(self, event, x, y, flags, param): 172 | x, y = np.int16([x, y]) # BUG 173 | if event == cv.EVENT_LBUTTONDOWN: 174 | self.drag_start = (x, y) 175 | return 176 | if self.drag_start: 177 | if flags & cv.EVENT_FLAG_LBUTTON: 178 | xo, yo = self.drag_start 179 | x0, y0 = np.minimum([xo, yo], [x, y]) 180 | x1, y1 = np.maximum([xo, yo], [x, y]) 181 | self.drag_rect = None 182 | if x1-x0 > 0 and y1-y0 > 0: 183 | self.drag_rect = (x0, y0, x1, y1) 184 | else: 185 | rect = self.drag_rect 186 | self.drag_start = None 187 | self.drag_rect = None 188 | if rect: 189 | self.callback(rect) 190 | def draw(self, vis): 191 | if not self.drag_rect: 192 | return False 193 | x0, y0, x1, y1 = self.drag_rect 194 | cv.rectangle(vis, (x0, y0), (x1, y1), (0, 255, 0), 2) 195 | return True 196 | @property 197 | def dragging(self): 198 | return self.drag_rect is not None 199 | 200 | 201 | def grouper(n, iterable, fillvalue=None): 202 | '''grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx''' 203 | args = [iter(iterable)] * n 204 | if PY3: 205 | output = it.zip_longest(fillvalue=fillvalue, *args) 206 | else: 207 | output = it.izip_longest(fillvalue=fillvalue, *args) 208 | return output 209 | 210 | def mosaic(w, imgs): 211 | '''Make a grid from images. 212 | w -- number of grid columns 213 | imgs -- images (must have same size and format) 214 | ''' 215 | imgs = iter(imgs) 216 | if PY3: 217 | img0 = next(imgs) 218 | else: 219 | img0 = imgs.next() 220 | pad = np.zeros_like(img0) 221 | imgs = it.chain([img0], imgs) 222 | rows = grouper(w, imgs, pad) 223 | return np.vstack(map(np.hstack, rows)) 224 | 225 | def getsize(img): 226 | h, w = img.shape[:2] 227 | return w, h 228 | 229 | def mdot(*args): 230 | return reduce(np.dot, args) 231 | 232 | def draw_keypoints(vis, keypoints, color=(0, 255, 255)): 233 | """Draw keypoints on an image.""" 234 | for kp in keypoints: 235 | x, y = kp.pt 236 | cv.circle(vis, (int(x), int(y)), 2, color) 237 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | """Vehicle speed estimation using the Lucas-Kanade tracker. 4 | 5 | This script tracks feature points across frames to estimate the speed of 6 | vehicles appearing in a video stream. It can read frames either from a 7 | provided video file or from a webcam. 8 | 9 | Usage:: 10 | 11 | python main.py [] 12 | 13 | Press ``ESC`` to exit the application. 14 | """ 15 | 16 | from __future__ import print_function 17 | 18 | import sys 19 | import math 20 | import numpy as np 21 | import cv2 as cv 22 | import time 23 | 24 | import video 25 | from common import anorm2, draw_str 26 | 27 | class App: 28 | """ 29 | The main application class that runs the Lucas-Kanade tracker 30 | on the provided video source. 31 | """ 32 | def __init__(self, video_src): 33 | self.track_len = 2 34 | self.detect_interval = 4 35 | self.tracks = [] 36 | self.cam = video.create_capture(video_src) 37 | self.alpha = 0.5 38 | self.frame_idx = 0 39 | 40 | def run(self): 41 | """Execute the tracking loop and estimate vehicle speeds. 42 | 43 | Frames are processed one by one, feature points are tracked and the 44 | resulting speeds for each lane are displayed on the output video. 45 | """ 46 | # Lucas-Kanade parameters 47 | lk_params = dict(winSize=(15, 15), 48 | maxLevel=2, 49 | criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03)) 50 | 51 | feature_params = dict(maxCorners=500, 52 | qualityLevel=0.3, 53 | minDistance=7, 54 | blockSize=7) 55 | 56 | # Constants 57 | fps = 30 58 | px2m1 = 0.0895 59 | px2m2 = 0.088 60 | px2m3 = 0.0774 61 | px2m4 = 0.0767 62 | px2m5 = 0.0736 63 | ms2kmh = 3.6 64 | 65 | ret, first_frame = self.cam.read() 66 | cal_mask = np.zeros_like(first_frame[:, :, 0]) 67 | view_mask = np.zeros_like(first_frame[:, :, 0]) 68 | view_polygon = np.array([[440, 1920], [420, 220], [680, 250], [1080, 480], [1080, 1920]]) 69 | cal_polygon = np.array([[440, 600], [420, 350], [1080, 350], [1080, 600]]) 70 | prv1, prv2, prv3, prv4, prv5 = 0,0,0,0,0 71 | prn1, prn2, prn3, prn4, prn5 = 0, 0, 0, 0, 0 72 | ptn1, ptn2, ptn3, ptn4, ptn5 = 0, 0, 0, 0, 0 73 | 74 | polygon1 = np.array([[550, 490], [425, 500],[420, 570], [570, 570]]) 75 | polygon2 = np.array([[570, 570], [555, 490], [680, 480], [720, 564]]) 76 | polygon3 = np.array([[720, 564],[680, 480], [835, 470], [930, 540]]) 77 | polygon4 = np.array([[930, 550], [835, 470], [970, 470], [1060, 550]]) 78 | polygon5 = np.array([[1080, 550], [1070, 550],[970, 470], [1080, 470]]) 79 | 80 | cv.fillConvexPoly(cal_mask, cal_polygon, 1) 81 | cv.fillConvexPoly(view_mask, view_polygon, 1) 82 | 83 | fourcc = cv.VideoWriter_fourcc(*'XVID') 84 | out = cv.VideoWriter("output.mp4", fourcc, 30.0, (1080, 1920)) 85 | 86 | while (self.cam.isOpened()): 87 | _ret, frame = self.cam.read() 88 | if _ret: 89 | vis = frame.copy() 90 | cmask = frame.copy() 91 | 92 | mm1, mm2, mm3, mm4, mm5 = 0, 0, 0, 0, 0 93 | v1, v2, v3, v4, v5 = 0, 0, 0, 0, 0 94 | 95 | 96 | frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) 97 | frame_gray = cv.bitwise_and(frame_gray, frame_gray, mask=cal_mask) 98 | 99 | vis = cv.bitwise_and(vis, vis, mask=view_mask) 100 | cv.line(vis,(400, 575),(1080, 540),(0, 0, 255), 5) 101 | cv.line(vis, (400, 495), (1080, 460), (0, 0, 255), 5) 102 | 103 | cv.fillPoly(cmask, [polygon1], (120, 0, 120), cv.LINE_AA) 104 | cv.fillPoly(cmask, [polygon2], (120, 120, 0), cv.LINE_AA) 105 | cv.fillPoly(cmask, [polygon3], (0, 120, 120), cv.LINE_AA) 106 | cv.fillPoly(cmask, [polygon4], (80, 0, 255), cv.LINE_AA) 107 | cv.fillPoly(cmask, [polygon5], (255, 0, 80), cv.LINE_AA) 108 | 109 | draw_str(vis, (30, 40), '1-lane speed: %d km/h' % prv1) 110 | draw_str(vis, (30, 80), '2-lane speed: %d km/h' % prv2) 111 | draw_str(vis, (30, 120), '3-lane speed: %d km/h' % prv3) 112 | draw_str(vis, (30, 160), '4-lane speed: %d km/h' % prv4) 113 | draw_str(vis, (30, 200), '5-lane speed: %d km/h' % prv5) 114 | draw_str(vis, (900, 40), 'ptn1: %d' % prn1) 115 | draw_str(vis, (900, 80), 'ptn2: %d' % prn2) 116 | draw_str(vis, (900, 120), 'ptn3: %d' % prn3) 117 | draw_str(vis, (900, 160), 'ptn4: %d' % prn4) 118 | draw_str(vis, (900, 200), 'ptn5: %d' %prn5) 119 | 120 | if len(self.tracks) > 0: 121 | img0, img1 = self.prev_gray, frame_gray 122 | p0 = np.float32([tr[-1] for tr in self.tracks]).reshape(-1, 1, 2) 123 | p1, _st, _err = cv.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params) 124 | p0r, _st, _err = cv.calcOpticalFlowPyrLK(img1, img0, p1, None, **lk_params) 125 | d = abs(p0-p0r).reshape(-1, 2).max(-1) 126 | good = d < 1 127 | new_tracks = [] 128 | for tr, (x, y), good_flag in zip(self.tracks, p1.reshape(-1, 2), good): 129 | if not good_flag: 130 | continue 131 | tr.append((x, y)) 132 | if len(tr) > self.track_len: 133 | del tr[0] 134 | new_tracks.append(tr) 135 | cv.circle(vis, (x, y), 3, (0, 255, 0), -1) 136 | self.tracks = new_tracks 137 | 138 | ptn1 = ptn2 = ptn3 = ptn4 = ptn5 = 0 139 | start = time.time() # record start time for profiling 140 | for idx, tr in enumerate(self.tracks): 141 | results = [ 142 | cv.pointPolygonTest(polygon1, tr[0], True), 143 | cv.pointPolygonTest(polygon2, tr[0], True), 144 | cv.pointPolygonTest(polygon3, tr[0], True), 145 | cv.pointPolygonTest(polygon4, tr[0], True), 146 | cv.pointPolygonTest(polygon5, tr[0], True), 147 | ] 148 | 149 | dist = np.linalg.norm(np.subtract(tr[0], tr[1])) 150 | if results[0] > 0: 151 | ptn1 += 1 152 | mm1 += dist 153 | v1 = (mm1 / ptn1) * px2m1 * fps * ms2kmh 154 | if results[1] > 0: 155 | ptn2 += 1 156 | mm2 += dist 157 | v2 = (mm2 / ptn2) * px2m2 * fps * ms2kmh 158 | if results[2] > 0: 159 | ptn3 += 1 160 | mm3 += dist 161 | v3 = (mm3 / ptn3) * px2m3 * fps * ms2kmh 162 | if results[3] > 0: 163 | ptn4 += 1 164 | mm4 += dist 165 | v4 = (mm4 / ptn4) * px2m4 * fps * ms2kmh 166 | if results[4] > 0: 167 | ptn5 += 1 168 | mm5 += dist 169 | v5 = (mm5 / ptn5) * px2m5 * fps * ms2kmh 170 | # print elapsed time for debugging 171 | # print("time :", time.time() - start) 172 | cv.polylines(vis, [np.int32(tr) for tr in self.tracks], False, (0, 0, 255)) 173 | 174 | prn1 = ptn1 175 | prn2 = ptn2 176 | prn3 = ptn3 177 | prn4 = ptn4 178 | prn5 = ptn5 179 | 180 | if self.frame_idx % self.detect_interval == 0: 181 | if ptn1 > 10: 182 | draw_str(vis, (900, 40), 'ptn1: %d' % ptn1) 183 | draw_str(vis, (30, 40), '1-lane speed: %d km/h' % v1) 184 | if ptn2 > 10: 185 | draw_str(vis, (900, 80), 'ptn2: %d' % ptn2) 186 | draw_str(vis, (30, 80), '2-lane speed: %d km/h' % v2) 187 | if ptn3 > 10: 188 | draw_str(vis, (900, 120), 'ptn3: %d' % ptn3) 189 | draw_str(vis, (30, 120), '3-lane speed: %d km/h' % v3) 190 | if ptn4 > 10: 191 | draw_str(vis, (900, 160), 'ptn4: %d' % ptn4) 192 | draw_str(vis, (30, 160), '4-lane speed: %d km/h' % v4) 193 | if ptn5 > 10: 194 | draw_str(vis, (900, 200), 'ptn5: %d' % ptn5) 195 | draw_str(vis, (30, 200), '5-lane speed: %d km/h' % v5) 196 | 197 | # Speed writing part 198 | prv1 = v1 199 | prv2 = v2 200 | prv3 = v3 201 | prv4 = v4 202 | prv5 = v5 203 | 204 | mask = np.zeros_like(frame_gray) 205 | mask[:] = 255 206 | for x, y in [np.int32(tr[-1]) for tr in self.tracks]: 207 | cv.circle(mask, (x, y), 3, 0, -1) 208 | p = cv.goodFeaturesToTrack(frame_gray, mask = mask, **feature_params) 209 | if p is not None: 210 | for x, y in np.float32(p).reshape(-1, 2): 211 | self.tracks.append([(x, y)])# 212 | 213 | 214 | self.frame_idx += 1 215 | self.prev_gray = frame_gray 216 | cv.addWeighted(cmask, self.alpha, vis, 1 - self.alpha, 0, vis) 217 | out.write(vis) 218 | #cv.imshow('lk_track', vis) 219 | #cv.waitKey(0) 220 | if cv.waitKey(1) & 0xFF == ord('q'): 221 | break 222 | else: 223 | break 224 | out.release() 225 | self.cam.release() 226 | cv.destroyAllWindows() 227 | 228 | def main(): 229 | """ 230 | The main entry point of the application. 231 | """ 232 | try: 233 | video_src = sys.argv[1] 234 | except IndexError: 235 | video_src = 0 236 | 237 | app = App(video_src) 238 | app.run() 239 | print('Done') 240 | 241 | 242 | if __name__ == '__main__': 243 | print(__doc__) 244 | main() 245 | -------------------------------------------------------------------------------- /tests/test_common.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | try: 3 | import numpy as np 4 | NUMPY_AVAILABLE = True 5 | except ImportError: # pragma: no cover - environment without numpy 6 | NUMPY_AVAILABLE = False 7 | np = None 8 | 9 | if NUMPY_AVAILABLE: 10 | import common 11 | 12 | @unittest.skipUnless(NUMPY_AVAILABLE, "numpy is required") 13 | class TestCommonFunctions(unittest.TestCase): 14 | def test_anorm(self): 15 | vec = np.array([3.0, 4.0]) 16 | self.assertAlmostEqual(common.anorm(vec), 5.0) 17 | self.assertAlmostEqual(common.anorm2(vec), 25.0) 18 | 19 | def test_mdot(self): 20 | a = np.eye(2) 21 | b = np.array([[2, 0], [0, 2]]) 22 | result = common.mdot(a, b) 23 | np.testing.assert_array_equal(result, b) 24 | 25 | def test_rect2rect_mtx(self): 26 | src = (0, 0, 1, 1) 27 | dst = (0, 0, 2, 2) 28 | M = common.rect2rect_mtx(src, dst) 29 | expected = np.array([[2.0, 0.0, 0.0], 30 | [0.0, 2.0, 0.0], 31 | [0.0, 0.0, 1.0]]) 32 | np.testing.assert_array_almost_equal(M, expected) 33 | 34 | if __name__ == '__main__': 35 | unittest.main() 36 | -------------------------------------------------------------------------------- /tst_scene_render.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | 4 | # Python 2/3 compatibility 5 | from __future__ import print_function 6 | 7 | import numpy as np 8 | import cv2 as cv 9 | 10 | from numpy import pi, sin, cos 11 | 12 | 13 | defaultSize = 512 14 | 15 | class TestSceneRender(): 16 | 17 | def __init__(self, bgImg = None, fgImg = None, 18 | deformation = False, speed = 0.25, **params): 19 | self.time = 0.0 20 | self.timeStep = 1.0 / 30.0 21 | self.foreground = fgImg 22 | self.deformation = deformation 23 | self.speed = speed 24 | 25 | if bgImg is not None: 26 | self.sceneBg = bgImg.copy() 27 | else: 28 | self.sceneBg = np.zeros(defaultSize, defaultSize, np.uint8) 29 | 30 | self.w = self.sceneBg.shape[0] 31 | self.h = self.sceneBg.shape[1] 32 | 33 | if fgImg is not None: 34 | self.foreground = fgImg.copy() 35 | self.center = self.currentCenter = (int(self.w/2 - fgImg.shape[0]/2), int(self.h/2 - fgImg.shape[1]/2)) 36 | 37 | self.xAmpl = self.sceneBg.shape[0] - (self.center[0] + fgImg.shape[0]) 38 | self.yAmpl = self.sceneBg.shape[1] - (self.center[1] + fgImg.shape[1]) 39 | 40 | self.initialRect = np.array([ (self.h/2, self.w/2), (self.h/2, self.w/2 + self.w/10), 41 | (self.h/2 + self.h/10, self.w/2 + self.w/10), (self.h/2 + self.h/10, self.w/2)]).astype(int) 42 | self.currentRect = self.initialRect 43 | 44 | def getXOffset(self, time): 45 | return int( self.xAmpl*cos(time*self.speed)) 46 | 47 | 48 | def getYOffset(self, time): 49 | return int(self.yAmpl*sin(time*self.speed)) 50 | 51 | def setInitialRect(self, rect): 52 | self.initialRect = rect 53 | 54 | def getRectInTime(self, time): 55 | 56 | if self.foreground is not None: 57 | tmp = np.array(self.center) + np.array((self.getXOffset(time), self.getYOffset(time))) 58 | x0, y0 = tmp 59 | x1, y1 = tmp + self.foreground.shape[0:2] 60 | return np.array([y0, x0, y1, x1]) 61 | else: 62 | x0, y0 = self.initialRect[0] + np.array((self.getXOffset(time), self.getYOffset(time))) 63 | x1, y1 = self.initialRect[2] + np.array((self.getXOffset(time), self.getYOffset(time))) 64 | return np.array([y0, x0, y1, x1]) 65 | 66 | def getCurrentRect(self): 67 | 68 | if self.foreground is not None: 69 | 70 | x0 = self.currentCenter[0] 71 | y0 = self.currentCenter[1] 72 | x1 = self.currentCenter[0] + self.foreground.shape[0] 73 | y1 = self.currentCenter[1] + self.foreground.shape[1] 74 | return np.array([y0, x0, y1, x1]) 75 | else: 76 | x0, y0 = self.currentRect[0] 77 | x1, y1 = self.currentRect[2] 78 | return np.array([x0, y0, x1, y1]) 79 | 80 | def getNextFrame(self): 81 | img = self.sceneBg.copy() 82 | 83 | if self.foreground is not None: 84 | self.currentCenter = (self.center[0] + self.getXOffset(self.time), self.center[1] + self.getYOffset(self.time)) 85 | img[self.currentCenter[0]:self.currentCenter[0]+self.foreground.shape[0], 86 | self.currentCenter[1]:self.currentCenter[1]+self.foreground.shape[1]] = self.foreground 87 | else: 88 | self.currentRect = self.initialRect + np.int( 30*cos(self.time*self.speed) + 50*sin(self.time*self.speed)) 89 | if self.deformation: 90 | self.currentRect[1:3] += int(self.h/20*cos(self.time)) 91 | cv.fillConvexPoly(img, self.currentRect, (0, 0, 255)) 92 | 93 | self.time += self.timeStep 94 | return img 95 | 96 | def resetTime(self): 97 | self.time = 0.0 98 | 99 | 100 | def main(): 101 | backGr = cv.imread(cv.samples.findFile('graf1.png')) 102 | fgr = cv.imread(cv.samples.findFile('box.png')) 103 | 104 | render = TestSceneRender(backGr, fgr) 105 | 106 | while True: 107 | 108 | img = render.getNextFrame() 109 | cv.imshow('img', img) 110 | 111 | ch = cv.waitKey(3) 112 | if ch == 27: 113 | break 114 | 115 | print('Done') 116 | 117 | 118 | if __name__ == '__main__': 119 | print(__doc__) 120 | main() 121 | cv.destroyAllWindows() 122 | -------------------------------------------------------------------------------- /video.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | ''' 4 | Video capture sample. 5 | Sample shows how VideoCapture class can be used to acquire video 6 | frames from a camera of a movie file. Also the sample provides 7 | an example of procedural video generation by an object, mimicking 8 | the VideoCapture interface (see Chess class). 9 | 'create_capture' is a convenience function for capture creation, 10 | falling back to procedural video in case of error. 11 | Usage: 12 | video.py [--shotdir ] [source0] [source1] ...' 13 | sourceN is an 14 | - integer number for camera capture 15 | - name of video file 16 | - synth: for procedural video 17 | Synth examples: 18 | synth:bg=lena.jpg:noise=0.1 19 | synth:class=chess:bg=lena.jpg:noise=0.1:size=640x480 20 | Keys: 21 | ESC - exit 22 | SPACE - save current frame to directory 23 | ''' 24 | 25 | # Python 2/3 compatibility 26 | from __future__ import print_function 27 | 28 | import numpy as np 29 | import cv2 as cv 30 | 31 | import re 32 | 33 | from numpy import pi, sin, cos 34 | 35 | # local modules 36 | from tst_scene_render import TestSceneRender 37 | import common 38 | 39 | class VideoSynthBase(object): 40 | def __init__(self, size=None, noise=0.0, bg = None, **params): 41 | self.bg = None 42 | self.frame_size = (640, 480) 43 | if bg is not None: 44 | self.bg = cv.imread(cv.samples.findFile(bg)) 45 | h, w = self.bg.shape[:2] 46 | self.frame_size = (w, h) 47 | 48 | if size is not None: 49 | w, h = map(int, size.split('x')) 50 | self.frame_size = (w, h) 51 | self.bg = cv.resize(self.bg, self.frame_size) 52 | 53 | self.noise = float(noise) 54 | 55 | def render(self, dst): 56 | pass 57 | 58 | def read(self, dst=None): 59 | w, h = self.frame_size 60 | 61 | if self.bg is None: 62 | buf = np.zeros((h, w, 3), np.uint8) 63 | else: 64 | buf = self.bg.copy() 65 | 66 | self.render(buf) 67 | 68 | if self.noise > 0.0: 69 | noise = np.zeros((h, w, 3), np.int8) 70 | cv.randn(noise, np.zeros(3), np.ones(3)*255*self.noise) 71 | buf = cv.add(buf, noise, dtype=cv.CV_8UC3) 72 | return True, buf 73 | 74 | def isOpened(self): 75 | return True 76 | 77 | class Book(VideoSynthBase): 78 | def __init__(self, **kw): 79 | super(Book, self).__init__(**kw) 80 | backGr = cv.imread(cv.samples.findFile('graf1.png')) 81 | fgr = cv.imread(cv.samples.findFile('box.png')) 82 | self.render = TestSceneRender(backGr, fgr, speed = 1) 83 | 84 | def read(self, dst=None): 85 | noise = np.zeros(self.render.sceneBg.shape, np.int8) 86 | cv.randn(noise, np.zeros(3), np.ones(3)*255*self.noise) 87 | 88 | return True, cv.add(self.render.getNextFrame(), noise, dtype=cv.CV_8UC3) 89 | 90 | class Cube(VideoSynthBase): 91 | def __init__(self, **kw): 92 | super(Cube, self).__init__(**kw) 93 | self.render = TestSceneRender(cv.imread(cv.samples.findFile('pca_test1.jpg')), deformation = True, speed = 1) 94 | 95 | def read(self, dst=None): 96 | noise = np.zeros(self.render.sceneBg.shape, np.int8) 97 | cv.randn(noise, np.zeros(3), np.ones(3)*255*self.noise) 98 | 99 | return True, cv.add(self.render.getNextFrame(), noise, dtype=cv.CV_8UC3) 100 | 101 | class Chess(VideoSynthBase): 102 | def __init__(self, **kw): 103 | super(Chess, self).__init__(**kw) 104 | 105 | w, h = self.frame_size 106 | 107 | self.grid_size = sx, sy = 10, 7 108 | white_quads = [] 109 | black_quads = [] 110 | for i, j in np.ndindex(sy, sx): 111 | q = [[j, i, 0], [j+1, i, 0], [j+1, i+1, 0], [j, i+1, 0]] 112 | [white_quads, black_quads][(i + j) % 2].append(q) 113 | self.white_quads = np.float32(white_quads) 114 | self.black_quads = np.float32(black_quads) 115 | 116 | fx = 0.9 117 | self.K = np.float64([[fx*w, 0, 0.5*(w-1)], 118 | [0, fx*w, 0.5*(h-1)], 119 | [0.0,0.0, 1.0]]) 120 | 121 | self.dist_coef = np.float64([-0.2, 0.1, 0, 0]) 122 | self.t = 0 123 | 124 | def draw_quads(self, img, quads, color = (0, 255, 0)): 125 | img_quads = cv.projectPoints(quads.reshape(-1, 3), self.rvec, self.tvec, self.K, self.dist_coef) [0] 126 | img_quads.shape = quads.shape[:2] + (2,) 127 | for q in img_quads: 128 | cv.fillConvexPoly(img, np.int32(q*4), color, cv.LINE_AA, shift=2) 129 | 130 | def render(self, dst): 131 | t = self.t 132 | self.t += 1.0/30.0 133 | 134 | sx, sy = self.grid_size 135 | center = np.array([0.5*sx, 0.5*sy, 0.0]) 136 | phi = pi/3 + sin(t*3)*pi/8 137 | c, s = cos(phi), sin(phi) 138 | ofs = np.array([sin(1.2*t), cos(1.8*t), 0]) * sx * 0.2 139 | eye_pos = center + np.array([cos(t)*c, sin(t)*c, s]) * 15.0 + ofs 140 | target_pos = center + ofs 141 | 142 | R, self.tvec = common.lookat(eye_pos, target_pos) 143 | self.rvec = common.mtx2rvec(R) 144 | 145 | self.draw_quads(dst, self.white_quads, (245, 245, 245)) 146 | self.draw_quads(dst, self.black_quads, (10, 10, 10)) 147 | 148 | 149 | classes = dict(chess=Chess, book=Book, cube=Cube) 150 | 151 | presets = dict( 152 | empty = 'synth:', 153 | lena = 'synth:bg=lena.jpg:noise=0.1', 154 | chess = 'synth:class=chess:bg=lena.jpg:noise=0.1:size=640x480', 155 | book = 'synth:class=book:bg=graf1.png:noise=0.1:size=640x480', 156 | cube = 'synth:class=cube:bg=pca_test1.jpg:noise=0.0:size=640x480' 157 | ) 158 | 159 | 160 | def create_capture(source = 0, fallback = presets['chess']): 161 | '''source: or '||synth [:= [:...]]' 162 | ''' 163 | source = str(source).strip() 164 | 165 | # Win32: handle drive letter ('c:', ...) 166 | source = re.sub(r'(^|=)([a-zA-Z]):([/\\a-zA-Z0-9])', r'\1?disk\2?\3', source) 167 | chunks = source.split(':') 168 | chunks = [re.sub(r'\?disk([a-zA-Z])\?', r'\1:', s) for s in chunks] 169 | 170 | source = chunks[0] 171 | try: source = int(source) 172 | except ValueError: pass 173 | params = dict( s.split('=') for s in chunks[1:] ) 174 | 175 | cap = None 176 | if source == 'synth': 177 | Class = classes.get(params.get('class', None), VideoSynthBase) 178 | try: cap = Class(**params) 179 | except: pass 180 | else: 181 | cap = cv.VideoCapture(source) 182 | if 'size' in params: 183 | w, h = map(int, params['size'].split('x')) 184 | cap.set(cv.CAP_PROP_FRAME_WIDTH, w) 185 | cap.set(cv.CAP_PROP_FRAME_HEIGHT, h) 186 | if cap is None or not cap.isOpened(): 187 | print('Warning: unable to open video source: ', source) 188 | if fallback is not None: 189 | return create_capture(fallback, None) 190 | return cap 191 | 192 | if __name__ == '__main__': 193 | import sys 194 | import getopt 195 | 196 | print(__doc__) 197 | 198 | args, sources = getopt.getopt(sys.argv[1:], '', 'shotdir=') 199 | args = dict(args) 200 | shotdir = args.get('--shotdir', '.') 201 | if len(sources) == 0: 202 | sources = [ 0 ] 203 | 204 | caps = list(map(create_capture, sources)) 205 | shot_idx = 0 206 | while True: 207 | imgs = [] 208 | for i, cap in enumerate(caps): 209 | ret, img = cap.read() 210 | imgs.append(img) 211 | cv.imshow('capture %d' % i, img) 212 | ch = cv.waitKey(1) 213 | if ch == 27: 214 | break 215 | if ch == ord(' '): 216 | for i, img in enumerate(imgs): 217 | fn = '%s/shot_%d_%03d.bmp' % (shotdir, i, shot_idx) 218 | cv.imwrite(fn, img) 219 | print(fn, 'saved') 220 | shot_idx += 1 221 | cv.destroyAllWindows() --------------------------------------------------------------------------------