src/capture.py

   1 #!/usr/bin/env python
   2
   3 """Go image capture.
   4
   5 This module defines a UI for capturing images with a (web)camera and imediately
   6 proccessing them with Imago.
   7 """
   8
   9 import os
  10 import sys
  11 import argparse
  12 import time
  13 from threading import Thread
  14 from Queue import Queue, Empty
  15
  16 import pygame
  17 from pygame.locals import QUIT, KEYDOWN
  18 import Image
  19
  20 from camera import Camera
  21
  22 class Screen:
  23     """Basic PyGame setup."""
  24     def __init__(self, res):
  25         pygame.init()
  26         pygame.display.set_mode(res)
  27         pygame.display.set_caption("Go image capture")
  28         self._screen = pygame.display.get_surface()
  29
  30     def display_picture(self, im):
  31         """Display image on PyGame screen."""
  32         pg_img = pygame.image.frombuffer(im.tostring(), im.size, im.mode)
  33         self._screen.blit(pg_img, (0,0))
  34         pygame.display.flip()
  35
  36 class Capture:
  37     """This object maintains communication between the camera, the PyGame screen
  38     and Imago."""
  39     def __init__(self, device, res):
  40         self.cam = Camera(vid=device, res=res)
  41         self.screen = Screen(res)
  42
  43         self.im_number = 0
  44
  45         self.saving_dir = "./captured/" + time.strftime("%Y-%m-%d %H:%M/")
  46
  47         if not os.path.isdir(self.saving_dir):
  48             os.makedirs(self.saving_dir)
  49
  50     def __del__(self):
  51         del self.cam
  52
  53     def live(self, q):
  54         """Run live preview on the screen."""
  55         clock = pygame.time.Clock()
  56         done = False
  57         while not done:
  58             if q:
  59                 try:
  60                     line = q.get_nowait() # or q.get(timeout=.1)
  61                 except Empty:
  62                     pass
  63                 else:
  64                     if line == "stop\n":
  65                         done = True
  66                     elif line == "exit\n":
  67                         sys.exit()
  68             else:
  69                 for event in pygame.event.get():
  70                     if event.type == pygame.QUIT:
  71                         sys.exit()
  72                     elif event.type == pygame.KEYDOWN:
  73                         done = True
  74
  75             im = self.cam.get_image()
  76             self.screen.display_picture(im)
  77             clock.tick(5)
  78
  79     def auto(self, interval):
  80         """Take new image every *interval* seconds.""" #TODO or is it milisecs?
  81         last = 0
  82         clock = pygame.time.Clock()
  83
  84         done = False
  85         while not done:
  86             for event in pygame.event.get():
  87                 if event.type == pygame.QUIT:
  88                     done = True
  89             if time.time() - last > interval:
  90                 last = time.time()
  91                 self.take()
  92             clock.tick(15)
  93
  94     def manual(self):
  95         """Take images manualy by pressing a key."""
  96         while True:
  97             event = pygame.event.wait()
  98             if event.type == QUIT:
  99                 break
 100             if event.type != KEYDOWN:
 101                 continue
 102             self.take()
 103
 104     def take(self):
 105         """Take a new image from the camera."""
 106         print "taking pic"
 107         im = self.cam.get_image()
 108         self.screen.display_picture(im)
 109         im.save(self.saving_dir + "{0:0>3}.jpg".format(self.im_number), 'JPEG')
 110         self.im_number += 1
 111
 112
 113 def main():
 114     """Parse the argument and setup the UI."""
 115     parser = argparse.ArgumentParser(description=__doc__)
 116     parser.add_argument('-c', '--cmd', dest='cmd', action='store_true',
 117                     help="take commands from stdin")
 118     parser.add_argument('-d', type=int, default=0,
 119                         help="video device id")
 120     parser.add_argument('-a', type=int, default=0,
 121                         help="take picture automaticaly every A seconds")
 122     parser.add_argument('-r', type=int, nargs=2, default=[640, 480],
 123                         help="set camera resolution")
 124     args = parser.parse_args()
 125
 126     res=(args.r[0], args.r[1])
 127     capture = Capture(args.d, res)
 128
 129
 130     if args.cmd:
 131         def enqueue_input(queue):
 132             for line in iter(sys.stdin.readline, b''):
 133                 queue.put(line)
 134
 135         q = Queue()
 136         t = Thread(target=enqueue_input, args=(q,))
 137         t.daemon = True
 138         t.start()
 139
 140         capture.live(q)
 141
 142         clock = pygame.time.Clock()
 143
 144         while True:
 145             try:
 146                 line = q.get_nowait() # or q.get(timeout=.1)
 147             except Empty:
 148                 pass
 149             else:
 150                 if line == "take\n":
 151                     capture.take()
 152                 elif line == "exit\n":
 153                     break
 154             clock.tick(10)
 155
 156     else:
 157         capture.live(None)
 158         if args.a > 0:
 159             capture.auto(args.a)
 160         else:
 161             capture.manual()
 162
 163     del capture
 164
 165 if __name__ == '__main__':
 166     try:
 167         main()
 168     finally:
 169         pygame.quit()