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         while True:
  57             if q:
  58                 try:
  59                     line = q.get_nowait() # or q.get(timeout=.1)
  60                 except Empty:
  61                     pass
  62                 else:
  63                     if line == "stop\n":
  64                         done = True
  65                     elif line == "exit\n":
  66                         sys.exit()
  67             else:
  68                 for event in pygame.event.get():
  69                     if event.type == pygame.QUIT:
  70                         sys.exit()
  71                     elif event.type == pygame.KEYDOWN:
  72                         done = True
  73
  74             im = self.cam.get_image()
  75             self.screen.display_picture(im)
  76             clock.tick(5)
  77
  78     def auto(self, interval):
  79         """Take new image every *interval* seconds.""" #TODO or is it milisecs?
  80         last = 0
  81         clock = pygame.time.Clock()
  82
  83         while True:
  84             for event in pygame.event.get():
  85                 if event.type == pygame.QUIT:
  86                     done = True
  87             if time.time() - last > interval:
  88                 last = time.time()
  89                 self.take()
  90             clock.tick(15)
  91
  92     def manual(self):
  93         """Take images manualy by pressing a key."""
  94         while True:
  95             event = pygame.event.wait()
  96             if event.type == QUIT:
  97                 break
  98             if event.type != KEYDOWN:
  99                 continue
 100             self.take()
 101
 102     def take(self):
 103         """Take a new image from the camera."""
 104         im = self.cam.get_image()
 105         self.screen.display_picture(im)
 106         im.save(self.saving_dir + "{0:0>3}.jpg".format(self.im_number), 'JPEG')
 107         self.im_number += 1
 108
 109
 110 def main():
 111     """Parse the argument and setup the UI."""
 112     parser = argparse.ArgumentParser(description=__doc__)
 113     parser.add_argument('-c', '--cmd', dest='cmd', action='store_true',
 114                     help="take commands from stdin")
 115     parser.add_argument('-d', type=int, default=0,
 116                         help="video device id")
 117     parser.add_argument('-a', type=int, default=0,
 118                         help="take picture automaticaly every A seconds")
 119     parser.add_argument('-r', type=int, nargs=2, default=[640, 480],
 120                         help="set camera resolution")
 121     args = parser.parse_args()
 122
 123     res=(args.r[0], args.r[1])
 124     capture = Capture(args.d, res)
 125
 126
 127     if args.cmd:
 128         def enqueue_input(queue):
 129             for line in iter(sys.stdin.readline, b''):
 130                 queue.put(line)
 131
 132         q = Queue()
 133         t = Thread(target=enqueue_input, args=(q,))
 134         t.daemon = True
 135         t.start()
 136
 137         capture.live(q)
 138
 139         clock = pygame.time.Clock()
 140
 141         while True:
 142             try:
 143                 line = q.get_nowait() # or q.get(timeout=.1)
 144             except Empty:
 145                 pass
 146             else:
 147                 if line == "take\n":
 148                     capture.take()
 149                 elif line == "exit\n":
 150                     break
 151             clock.tick(10)
 152
 153     else:
 154         capture.live(None)
 155         if args.a > 0:
 156             capture.auto(args.a)
 157         else:
 158             capture.manual()
 159
 160     del capture
 161
 162 if __name__ == '__main__':
 163     try:
 164         main()
 165     finally:
 166         pygame.quit()