OpenSeeFace/input_reader.py

import re
import sys
import os
import cv2
import numpy as np
import dshowcapture

class VideoReader():
    def __init__(self, capture):
        if os.name == 'nt':
            self.cap = cv2.VideoCapture(capture, cv2.CAP_DSHOW)
        else:
            self.cap = cv2.VideoCapture(capture)
        if self.cap is None:
            print("The video source cannot be opened")
            sys.exit(0)
    def is_open(self):
        return self.cap.isOpened()
    def is_ready(self):
        return True
    def read(self):
        return self.cap.read()
    def close(self):
        self.cap.release()

class CameraReader(VideoReader):
    def __init__(self, capture, width, height, fps, use_dshowcapture=True):
        self.device = None
        self.width = width
        self.height = height
        self.fps = fps
        if os.name == 'nt' and use_dshowcapture:
            self.device = dshowcapture.DShowCapture()
            self.device.capture_device(capture, self.width, self.height, self.fps)
        else:
            super(CameraReader, self).__init__(capture)
            self.cap.set(3, width)
            self.cap.set(4, height)
    def is_open(self):
        if self.device is None:
            return super(CameraReader, self).is_open()
        return self.device.capturing()
    def is_ready(self):
        if self.device is None:
            return super(CameraReader, self).is_ready()
        return self.device.capturing()
    def read(self):
        if self.device is None:
            return super(CameraReader, self).read()
        img = self.device.get_frame(1000)
        if img is None:
            return False, None
        else:
            return True, img
    def close(self):
        if self.device is None:
            super(CameraReader, self).close()
        else:
            self.device.destroy_capture()

class RawReader:
    def __init__(self, width, height):
        self.width = int(width)
        self.height = int(height)
        
        if self.width < 1 or self.height < 1:
            print("No acceptable size was given for reading raw RGB frames.")
            sys.exit(0)

        self.len = self.width * self.height * 3
        self.open = True
    def is_open(self):
        return self.open
    def is_ready(self):
        return True
    def read(self):
        frame = bytearray()
        read_bytes = 0
        while read_bytes < self.len:
            bytes = sys.stdin.buffer.read(self.len)
            read_bytes += len(bytes)
            frame.extend(bytes)
        return True, np.frombuffer(frame, dtype=np.uint8).reshape((self.height, self.width, 3))
    def close(self):
        self.open = False

def try_int(s):
    try:
        return int(s)
    except:
        return None

class InputReader():
    def __init__(self, capture, raw_rgb, width, height, fps, use_dshowcapture=False):
        self.reader = None
        try:
            if raw_rgb > 0:
                self.reader = RawReader(width, height)
            elif os.path.exists(capture):
                self.reader = VideoReader(capture)
            elif capture == str(try_int(capture)):
                self.reader = CameraReader(int(capture), width, height, fps, use_dshowcapture=use_dshowcapture)
        except Exception as e:
            print("Error: " + str(e))

        if self.reader is None or not self.reader.is_open():
            print("There was no valid input.")
            sys.exit(0)
    def is_open(self):
        return self.reader.is_open()
    def is_ready(self):
        return self.reader.is_ready()
    def read(self):
        return self.reader.read()
    def close(self):
        self.reader.close()