shepherd-agent/shepherd/plugins/usbcam.py

import io
import os
from datetime import datetime
import time
import re

import shepherd.config as shconf
import shepherd.plugin

import threading

import subprocess

from collections import namedtuple, OrderedDict
from operator import itemgetter


import cv2
from PIL import Image, ImageDraw, ImageFont

asset_dir = os.path.dirname(os.path.realpath(__file__))

overlayfont_filename = os.path.join(asset_dir, "DejaVuSansMono.ttf")
logo_filename = os.path.join(asset_dir, "smallshepherd.png")

# Note: Add a lock to the gstreamer function, to avoid multiple triggers colliding

CameraPort = namedtuple(
    'CameraPort', ['usbPath', 'devicePath'])

# Short wrapper to allow use in a ``with`` context


class VideoCaptureCtx():
    def __init__(self, *args, **kwargs):
        self.capture_dev = cv2.VideoCapture(*args, **kwargs)

    def __enter__(self):
        return self.capture_dev

    def __exit__(self, *args):
        self.capture_dev.release()


def get_connected_cameras():
    # This will return devices orderd by the USB path, regardless of the order they're connected in
    device_list_str = subprocess.run(
        ['v4l2-ctl', '--list-devices'], text=True, stdout=subprocess.PIPE).stdout
    # in each match, first group is the USB path, second group is the device path
    portlist = re.findall(r"-([\d.]+?)\):\n\s*?(\/dev\S+?)\n", device_list_str)
    return [CameraPort(*port) for port in portlist]


def get_capture_formats(video_device):
    """
    Call ``v4l2-ctl --device {video_device} --list-formats-ext`` and parse the output into a format dict

    Returns a dict with 4CC format codes as keys, and lists of (width,height) tuples as values
    """
    device_fmt_str = subprocess.run(
        ['v4l2-ctl', '--device', F'{video_device}', '--list-formats-ext'], text=True, stdout=subprocess.PIPE).stdout

    split_fmts = re.split(r"\[\d\]: '(\w{4}).*", device_fmt_str)
    if len(split_fmts) < 3:
        raise Exception("Did not get valid device format list output")

    # Iterate through successive pairs in the split, where the first is the format mode and the
    # second is the text containing all the resolution options. Skip the first bit, which is rubbish
    format_dict = {}
    for fourcc, size_text in zip(split_fmts[1::2], split_fmts[2::2]):
        resolutions = re.findall(r"(\d+?)x(\d+?)\D", size_text)
        format_dict[fourcc] = resolutions
    return format_dict


def get_largest_resolution(size_list):
    """
    Accepts a list of tuples where the first element is a width and the second is a height.

    Returns a single resolution tuple representing the largest area from the list
    """
    return max(size_list, key=lambda size: int(size[0]*int(size[1])))


def set_camera_format_v4l2(video_device, fourcc, width, height):
    """
    Set the camera device capture format using the external v4l2-ctl tool
    """
    subprocess.run(['v4l2-ctl', '--device', F'{video_device}',
                    F'--set-fmt-video width={width},height={height},pixelformat={fourcc}'], text=True)


def set_camera_format_opencv(capture_device, fourcc, width, height):
    """
    Set the camera device capture format using internal OpenCV set methods
    """
    # VideoWriter_fourcc expects a list of characters, so need to unpack the string
    capture_device.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*fourcc))
    capture_device.set(cv2.CAP_PROP_FRAME_WIDTH, int(width))
    capture_device.set(cv2.CAP_PROP_FRAME_HEIGHT, int(height))


class USBCamPlugin(shepherd.plugin.Plugin):
    @staticmethod
    def define_config(confdef):
        confdef.add_def('upload_images', shconf.BoolDef(default=False, optional=True,
                                                        helptext="If true, move to an Uploader bucket. Requires Uploader plugin"))
        confdef.add_def('upload_bucket', shconf.StringDef(default="", optional=True,
                                                          helptext="Name of uploader bucket to shift images to."))
        confdef.add_def('save_directory', shconf.StringDef(default="", optional=True,
                                                           helptext="Name of directory path to save images. If empty, a 'usbcamera' directory under the Shepherd root dir will be used"))
        confdef.add_def('append_id', shconf.BoolDef(default=True, optional=True,
                                                    helptext="If true, add the system ID to the end of image filenames"))
        confdef.add_def('show_overlay', shconf.BoolDef(default=True, optional=True,
                                                       helptext="If true, add an overlay on each image with the system ID and date."))
        confdef.add_def('overlay_desc', shconf.StringDef(default="", optional=True,
                                                         helptext="Text to add to the overlay after the system ID and camera name"))
        confdef.add_def('jpeg_quality', shconf.IntDef(default=85, minval=60, maxval=95, optional=True,
                                                      helptext="JPEG quality to save with. Max of 95, passed directly to Pillow"))
        confdef.add_def('stabilise_delay', shconf.IntDef(default=5, minval=1, maxval=30, optional=True,
                                                         helptext="Number of seconds to wait after starting each camera for exposure and white balance to settle"))

        array = confdef.add_def('trigger', shconf.DictListDef(
            helptext="Array of triggers that will use all cameras"))
        array.add_def('hour', shconf.StringDef())
        array.add_def('minute', shconf.StringDef())
        array.add_def('second', shconf.StringDef(default="0", optional=True))

        camarray = confdef.add_def('camera', shconf.DictListDef(
            helptext="List of cameras to try and connect to. Multiple ports may be listed, and any not connected will be skipped on each trigger."))
        camarray.add_def('name', shconf.StringDef(default="", optional=False,
                                                  helptext="Name of camera, appended to filename and added to overlay"))
        camarray.add_def('usb_port', shconf.StringDef(default="*", optional=False,
                                                      helptext="USB port descriptor of the from '3.4.1' (which would indicate port1 on a hub plugged into port4 on a hub plugged into port 3 of the system). This can be found by running 'v4l2-ctl --list-devices'. A single camera with a wildcard '*' port is also allowed, and will match any remaining available camera."))

    def __init__(self, pluginInterface, config):
        super().__init__(pluginInterface, config)
        self.config = config
        self.interface = pluginInterface
        self.plugins = pluginInterface.other_plugins
        self.hooks = pluginInterface.hooks

        self.root_dir = os.path.expanduser(pluginInterface.coreconfig["root_dir"])
        self.id = pluginInterface.coreconfig["name"]

        self.interface.register_hook("pre_cam")
        self.interface.register_hook("post_cam")
        self.interface.register_function(self.camera_job)
        self.interface.register_function(self.run_cameras)
        # do some camera init stuff

        print("USBCamera config:")
        print(self.config)

        self.gstlock = threading.Lock()

        if self.config["save_directory"] == "":
            self.save_directory = os.path.join(self.root_dir, "usbcamera")
        else:
            self.save_directory = self.config["save_directory"]

        if not os.path.exists(self.save_directory):
            os.makedirs(self.save_directory)

        if self.config["show_overlay"]:
            # Load assets
            self.logo_im = Image.open(logo_filename)

            self.font_size_cache = {}
            self.logo_size_cache = {}

        # Dict of camera names storing the USB path as the value
        self.defined_cams = OrderedDict()
        # List of wildcard camera names
        self.wildcard_cams = []

        # Go through camera configs sorted by name
        for camera in sorted(self.config["camera"], key=itemgetter("name")):
            if camera["name"] in self.defined_cams:
                raise shconf.InvalidConfigError(
                    "Can't have more than one usb camera defined with the same config name")
            if camera["usb_port"] == '*':
                self.wildcard_cams.append(camera["name"])
            else:
                self.defined_cams[camera["name"]] = camera["usb_port"]

        for trigger in self.config["trigger"]:
            trigger_id = trigger["hour"]+'-' + trigger["minute"]+'-'+trigger["second"]
            self.interface.add_job(
                self.camera_job, trigger["hour"], trigger["minute"], trigger["second"], job_name=trigger_id)

    def _generate_overlay(self, width, height, image_time, camera_name):
        font_size = int(height/40)
        margin_size = int(font_size/5)

        if font_size not in self.font_size_cache:
            self.font_size_cache[font_size] = ImageFont.truetype(
                overlayfont_filename, int(font_size*0.9))
        thisfont = self.font_size_cache[font_size]

        if font_size not in self.logo_size_cache:
            newsize = (int(self.logo_im.width*(
                font_size/self.logo_im.height)), font_size)
            self.logo_size_cache[font_size] = self.logo_im.resize(
                newsize, Image.BILINEAR)
        thislogo = self.logo_size_cache[font_size]

        desc_text = camera_name + " " + self.config["overlay_desc"]
        if self.config["append_id"]:
            desc_text = self.id + "  " + desc_text

        time_text = image_time.strftime("%Y-%m-%d %H:%M:%S")

        overlay = Image.new('RGBA', (width, font_size+(2*margin_size)), (0, 0, 0))
        overlay.paste(thislogo, (int((overlay.width-thislogo.width)/2), margin_size))

        draw = ImageDraw.Draw(overlay)
        draw.text((margin_size*2, margin_size), desc_text,
                  font=thisfont, fill=(255, 255, 255, 255))

        datewidth, _ = draw.textsize(time_text, thisfont)
        draw.text((overlay.width-(margin_size*2)-datewidth, margin_size), time_text, font=thisfont,
                  fill=(255, 255, 255, 255))

        # make whole overlay half transparent
        overlay.putalpha(128)
        return overlay

    def _process_image(self, cv_frame, camera_name):

        image_time = datetime.now()

        # Convert over to PIL. Mostly so we can use our own font.
        img = Image.fromarray(cv2.cvtColor(cv_frame, cv2.COLOR_BGR2RGB))

        if self.config["show_overlay"]:
            overlay = self._generate_overlay(img.width, img.height, image_time, camera_name)
            img.paste(overlay, (0, img.height-overlay.height), overlay)

        image_filename = image_time.strftime("%Y-%m-%d %H-%M-%S")
        if self.config["append_id"]:
            image_filename = image_filename + " " + self.id

        if camera_name != "":
            image_filename = image_filename+" "+camera_name
        image_filename = image_filename + ".jpg"
        image_filename = os.path.join(self.save_directory, image_filename)
        img.save(image_filename+".writing", "JPEG", quality=self.config["jpeg_quality"])
        os.rename(image_filename+".writing", image_filename)

        if self.config["upload_images"]:
            self.plugins["uploader"].move_to_bucket(image_filename, self.config["upload_bucket"])

    def _capture_image(self, device_path, camera_name):
        print("Running camera "+camera_name)

        with self.gstlock:

            # gst_str = ('v4l2src device='+device_path+' ! '
            #           'videoconvert ! appsink drop=true max-buffers=1 sync=false')

            # vidcap = cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)

            fmts = get_capture_formats(device_path)

            with VideoCaptureCtx(device_path, cv2.CAP_V4L2) as vidcap:
                if "MJPG" in fmts:
                    size = get_largest_resolution(fmts["MJPG"])
                    set_camera_format_opencv(vidcap, "MJPG", size[0], size[1])

                # stream only starts after first grab

                print("Starting cam")
                read_flag, frame = vidcap.read()
                delay_start = time.time()
                while (time.time() - delay_start) < self.config["stabilise_delay"]:
                    vidcap.grab()
                # time.sleep(self.config["stabilise_delay"])
                # clear old buffer
                # print("Flushing capture")
                # vidcap.grab()
                print("Reading")
                read_flag, frame = vidcap.read()
                # print("Changing to YUYV")
                # if "YUYV" in fmts:
                #    size = get_largest_resolution(fmts["YUYV"])
                #    set_camera_format_opencv(vidcap, "YUYV", size[0], size[1])
                # print("Reading again")
                # read_flag, frame2 = vidcap.read()

        if read_flag:
            self._process_image(frame, camera_name)
            # self._process_image(frame2, camera_name+"(2)")
        else:
            print("Could not read camera "+camera_name +
                  " on USB port "+device_path)

    def run_cameras(self, name_suffix=""):
        connected_cams = OrderedDict(get_connected_cameras())

        for defined_name, defined_usb_path in self.defined_cams.items():
            if defined_usb_path in connected_cams:

                self._capture_image(connected_cams.pop(
                    defined_usb_path), defined_name+name_suffix)

            else:
                print("USB Camera "+defined_name+" on port " +
                      defined_usb_path+" is not currently connected")

        for cam_name in self.wildcard_cams:
            if len(connected_cams) > 0:
                self._capture_image(connected_cams.popitem(
                    last=False)[1], cam_name+name_suffix)
            else:
                print(
                    "No connected USB cameras are currently left to match to "+cam_name+" ")
                break

    def camera_job(self):
        self.hooks.pre_cam()
        self.run_cameras()
        self.hooks.post_cam()


if __name__ == "__main__":
    pass
    # print("main")
    # main(sys.argv[1:])