capture.py

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