input_fetcher.py

import os
import cv2
import re
import time
from threading import RLock
import numpy as np

from codependent_thread import CodependentThread
from image_misc import cv2_imshow_rgb, cv2_read_file_rgb, read_cam_frame, crop_to_square
from misc import tsplit


class InputImageFetcher(CodependentThread):
    '''Fetches images from a webcam or loads from a directory.'''
    
    def __init__(self, settings):
        CodependentThread.__init__(self, settings.input_updater_heartbeat_required)
        self.daemon = True
        self.lock = RLock()
        self.quit = False
        self.latest_frame_idx = -1
        self.latest_frame_data = None
        self.latest_frame_is_from_cam = False

        # True for loading from file, False for loading from camera
        self.static_file_mode = True
        self.settings = settings

        # True for streching the image, False for cropping largest square
        self.static_file_stretch_mode = self.settings.static_file_stretch_mode
        
        # Cam input
        self.capture_device = settings.input_updater_capture_device
        self.no_cam_present = (self.capture_device is None)     # Disable all cam functionality
        self.bound_cap_device = None
        self.sleep_after_read_frame = settings.input_updater_sleep_after_read_frame
        self.latest_cam_frame = None
        self.freeze_cam = False

        # Static file input

        # latest image filename selected, used to avoid reloading
        self.latest_static_filename = None

        # latest loaded image frame, holds the pixels and used to force reloading
        self.latest_static_frame = None

        # keeps current index of loaded file, doesn't seem important
        self.static_file_idx = None

        # contains the requested number of increaments for file index
        self.static_file_idx_increment = 0
        
    def bind_camera(self):
        # Due to OpenCV limitations, this should be called from the main thread
        print 'InputImageFetcher: bind_camera starting'
        if self.no_cam_present:
            print 'InputImageFetcher: skipping camera bind (device: None)'
        else:
            self.bound_cap_device = cv2.VideoCapture(self.capture_device)
            if self.bound_cap_device.isOpened():
                print 'InputImageFetcher: capture device %s is open' % self.capture_device
            else:
                print '\n\nWARNING: InputImageFetcher: capture device %s failed to open! Camera will not be available!\n\n' % self.capture_device
                self.bound_cap_device = None
                self.no_cam_present = True
        print 'InputImageFetcher: bind_camera finished'

    def free_camera(self):
        # Due to OpenCV limitations, this should be called from the main thread
        if self.no_cam_present:
            print 'InputImageFetcher: skipping camera free (device: None)'
        else:
            print 'InputImageFetcher: freeing camera'
            del self.bound_cap_device  # free the camera
            self.bound_cap_device = None
            print 'InputImageFetcher: camera freed'

    def set_mode_static(self):
        with self.lock:
            self.static_file_mode = True
        
    def set_mode_cam(self):
        with self.lock:
            if self.no_cam_present:
                print 'WARNING: ignoring set_mode_cam, no cam present'
            else:
                self.static_file_mode = False
                assert self.bound_cap_device != None, 'Call bind_camera first'
        
    def toggle_input_mode(self):
        with self.lock:
            if self.static_file_mode:
                self.set_mode_cam()
            else:
                self.set_mode_static()
        
    def set_mode_stretch_on(self):
        with self.lock:
            if not self.static_file_stretch_mode:
                self.static_file_stretch_mode = True
                self.latest_static_frame = None   # Force reload
                #self.latest_frame_is_from_cam = True  # Force reload
        
    def set_mode_stretch_off(self):
        with self.lock:
            if self.static_file_stretch_mode:
                self.static_file_stretch_mode = False
                self.latest_static_frame = None   # Force reload
                #self.latest_frame_is_from_cam = True  # Force reload
        
    def toggle_stretch_mode(self):
        with self.lock:
            if self.static_file_stretch_mode:
                self.set_mode_stretch_off()
            else:
                self.set_mode_stretch_on()
        
    def run(self):
        while not self.quit and not self.is_timed_out():
            #start_time = time.time()
            if self.static_file_mode:
                self.check_increment_and_load_image()
            else:
                if self.freeze_cam and self.latest_cam_frame is not None:
                    # If static file mode was switched to cam mode but cam is still frozen, we need to push the cam frame again
                    if not self.latest_frame_is_from_cam:
                        self._increment_and_set_frame(self.latest_cam_frame, True)
                else:
                    frame_full = read_cam_frame(self.bound_cap_device)
                    #print '====> just read frame', frame_full.shape
                    frame = crop_to_square(frame_full)
                    with self.lock:
                        self.latest_cam_frame = frame
                        self._increment_and_set_frame(self.latest_cam_frame, True)
            
            time.sleep(self.sleep_after_read_frame)
            #print 'Reading one frame took', time.time() - start_time

        print 'InputImageFetcher: exiting run method'
        #print 'InputImageFetcher: read', self.read_frames, 'frames'

    def get_frame(self):
        '''Fetch the latest frame_idx and frame. The idx increments
        any time the frame data changes. If the idx is < 0, the frame
        is not valid.
        '''
        with self.lock:
            return (self.latest_frame_idx, self.latest_frame_data)

    def increment_static_file_idx(self, amount = 1):
        with self.lock:
            self.static_file_idx_increment += amount

    def _increment_and_set_frame(self, frame, from_cam):
        assert frame is not None
        with self.lock:
            self.latest_frame_idx += 1
            self.latest_frame_data = frame
            self.latest_frame_is_from_cam = from_cam

    def get_files_from_directory(self):
        # returns list of files in requested directory

        available_files = []
        match_flags = re.IGNORECASE if self.settings.static_files_ignore_case else 0
        for filename in os.listdir(self.settings.static_files_dir):
            if re.match(self.settings.static_files_regexp, filename, match_flags):
                available_files.append(filename)

        return available_files

    def get_files_from_image_list(self):
        # returns list of files in requested image list file

        available_files = []

        with open(self.settings.static_files_input_file, 'r') as image_list_file:
            lines = image_list_file.readlines()
            # take first token from each line
            available_files = [tsplit(line, True,' ',',','\t')[0] for line in lines if line.strip() != ""]

        return available_files

    def get_files_from_siamese_image_list(self):
        # returns list of pair files in requested siamese image list file

        available_files = []

        with open(self.settings.static_files_input_file, 'r') as image_list_file:
            lines = image_list_file.readlines()
            # take first and second tokens from each line
            available_files = [(tsplit(line, True, ' ', ',','\t')[0], tsplit(line, True, ' ', ',','\t')[1])
                               for line in lines if line.strip() != ""]

        return available_files

    def check_increment_and_load_image(self):
        with self.lock:
            if (self.static_file_idx_increment == 0 and
                self.static_file_idx is not None and
                not self.latest_frame_is_from_cam and
                self.latest_static_frame is not None):
                # Skip if a static frame is already loaded and there is no increment
                return

            # available_files - local list of files
            if self.settings.static_files_input_mode == "directory":
                available_files = self.get_files_from_directory()
            elif self.settings.static_files_input_mode == "image_list":
                available_files = self.get_files_from_image_list()
            elif self.settings.static_files_input_mode == "siamese_image_list":
                available_files = self.get_files_from_siamese_image_list()
            else:
                raise Exception(('Error: setting static_files_input_mode has invalid option (%s)' %
                                (self.settings.static_files_input_mode) ))

            #print 'Found files:'
            #for filename in available_files:
            #    print '   %s' % filename
            assert len(available_files) != 0, ('Error: No files found in %s matching %s (current working directory is %s)' %
                                               (self.settings.static_files_dir, self.settings.static_files_regexp, os.getcwd()))
            if self.static_file_idx is None:
                self.static_file_idx = 0
            self.static_file_idx = (self.static_file_idx + self.static_file_idx_increment) % len(available_files)
            self.static_file_idx_increment = 0
            if self.latest_static_filename != available_files[self.static_file_idx] or self.latest_static_frame is None:
                self.latest_static_filename = available_files[self.static_file_idx]

                if self.settings.static_files_input_mode == "siamese_image_list":
                    # loading two images for siamese network
                    im1 = cv2_read_file_rgb(os.path.join(self.settings.static_files_dir, self.latest_static_filename[0]))
                    im2 = cv2_read_file_rgb(os.path.join(self.settings.static_files_dir, self.latest_static_filename[1]))
                    if not self.static_file_stretch_mode:
                        im1 = crop_to_square(im1)
                        im2 = crop_to_square(im2)

                    im = (im1,im2)

                else:
                    im = cv2_read_file_rgb(os.path.join(self.settings.static_files_dir, self.latest_static_filename))
                    if not self.static_file_stretch_mode:
                        im = crop_to_square(im)

                self.latest_static_frame = im
            self._increment_and_set_frame(self.latest_static_frame, False)