video_test.py
1 |
import pygame |
---|---|
2 |
import pygame.camera |
3 |
from pygame.locals import * |
4 |
import numpy as np |
5 |
|
6 |
pygame.init() |
7 |
pygame.camera.init() |
8 |
|
9 |
|
10 |
class Capture(object): |
11 |
def __init__(self): |
12 |
self.size = (640,480) |
13 |
# create a display surface. standard pygame stuff
|
14 |
self.display = pygame.display.set_mode(self.size, 0) |
15 |
|
16 |
# this is the same as what we saw before
|
17 |
self.clist = pygame.camera.list_cameras()
|
18 |
if not self.clist: |
19 |
raise ValueError("Sorry, no cameras detected.") |
20 |
self.cam = pygame.camera.Camera(self.clist[1], self.size,"HSV") |
21 |
self.cam.start()
|
22 |
|
23 |
# create a surface to capture to. for performance purposes
|
24 |
# bit depth is the same as that of the display surface.
|
25 |
self.snapshot = pygame.surface.Surface(self.size, 0, self.display) |
26 |
|
27 |
def get_and_flip(self): |
28 |
# if you don't want to tie the framerate to the camera, you can check
|
29 |
# if the camera has an image ready. note that while this works
|
30 |
# on most cameras, some will never return true.
|
31 |
if self.cam.query_image(): |
32 |
self.snapshot = self.cam.get_image(self.snapshot) |
33 |
self.do_CV();
|
34 |
# blit it to the display surface. simple!
|
35 |
# self.display.blit(self.snapshot, (0,0))
|
36 |
pygame.display.flip() |
37 |
|
38 |
def main(self): |
39 |
going = True
|
40 |
while going:
|
41 |
events = pygame.event.get() |
42 |
for e in events: |
43 |
if e.type == QUIT or (e.type == KEYDOWN and e.key == K_ESCAPE): |
44 |
# close the camera safely
|
45 |
self.cam.stop()
|
46 |
going = False
|
47 |
|
48 |
self.get_and_flip()
|
49 |
|
50 |
def do_CV(self): |
51 |
im = pygame.surfarray.pixels3d(self.snapshot)
|
52 |
im = im/255.0;
|
53 |
# for finding green
|
54 |
im[:,:,0] = (im[:,:,0]>.25).astype(int) * (im[:,:,0]<.45).astype(int) |
55 |
im[:,:,1] = (im[:,:,1]>.20).astype(int) * (im[:,:,1]<.65).astype(int) |
56 |
im[:,:,2] = (im[:,:,2]>.1).astype(int) * (im[:,:,2]<.55).astype(int) |
57 |
im[:,:,0] = im[:,:,0] * im[:,:,1] * im[:,:,2]; |
58 |
im[:,:,1] = im[:,:,0] |
59 |
im[:,:,2] = im[:,:,0] |
60 |
|
61 |
im = np.floor(255*im);
|
62 |
im = im.astype(int)
|
63 |
pygame.surfarray.blit_array(self.display,im)
|
64 |
|
65 |
|
66 |
|
67 |
cap = Capture(); |
68 |
cap.main() |