camera-capture

#!/usr/bin/env ruby
require 'yaml'
require 'gphoto2'
require 'rmagick'
require 'optparse'
require 'pathname'
require 'tty-prompt'
require 'tty-spinner'

# Main Camera Capture App Class
#
class CameraCapture

    # Specific error when no camera exists.
    class NoCameraError < StandardError; end

    # the default v4l2 output device, written to by ffmpeg.
    DEFAULT_OUTPUT_DEVICE='/dev/video9'

    # the default config location
    DEFAULT_CONFIG_LOCATION=Pathname( '%s/.config/camera-capture.yml' % [ ENV[ 'HOME' ] ] )

    # default config
    DEFAULT_CONFIG = {
        output_device:               DEFAULT_OUTPUT_DEVICE,
        output_geometry_width:       1920,
        output_geometry_height:      1280,
        output_geometry_aspect:      Rational( 1920, 1280 ),
        output_geometry_keep_aspect: true,
        wait_for_camera:             false,
        capture_timeout:             60*5
    }


    # Initialize the app
    #
    def initialize( **options )
        Thread.abort_on_exception = false
        @prompt = TTY::Prompt.new( quiet: true, interrupt: :noop )
        @command_queue = Queue.new
        @result_queue = Queue.new

        self.load_config( options )
    end


    # the TTY::Prompt object
    attr_accessor :prompt

    # the selected GPhoto2::Camera
    attr_accessor :selected_camera

    # the handle for the video capture/gphoto2 thread
    attr_accessor :capture_thread

    # the handle for the intermediary/oob thread
    attr_accessor :interface_thread

    # the queue for sending data into the capture thread
    attr_accessor :command_queue

    # the queue for getting data out of the capture thread
    attr_accessor :result_queue

    # the config hash
    attr_accessor :config

    # the captured video geometry
    attr_accessor :capture_geometry_width
    attr_accessor :capture_geometry_height
    attr_accessor :capture_geometry_aspect


    # run the program
    #
    def run
        self.check_v4l2
        self.select_camera
        self.run_interface_thread
        self.menu_main
    end


    #########
    protected
    #########

    # load/update the default configuration from commandline args and
    # configuration parameters
    #
    def load_config( options )
        # attempt to open the config file, if it exists
        fl = options[ :config ] || DEFAULT_CONFIG_LOCATION
        fc = fl.read rescue nil
        begin
            yaml = YAML.load( fc )
            yaml.transform_keys!{|k| k.to_sym }
        rescue => e
            self.prompt.warn "Config parse error:"
            self.prompt.say e.message
            yaml = {}
        end

        # iterate through all valid config options, favoring cmdline options over config file,
        # and defaulting to the defaults set in the hash at the beginning of the class.
        @config = DEFAULT_CONFIG
        @config.each do |k, _v|
            next if k == :output_geometry_aspect
            @config[ k ] = options[ k ] || yaml[ k ] || @config[ k ]
        end

        # calculated values
        @config[ :output_geometry_aspect ] = Rational( @config[ :output_geometry_width ], @config[ :output_geometry_height ] )
    end


    ##############################################
    # C A M E R A  F U N C T I O N S
    ##############################################

    # find a camera to use
    #
    def select_camera
        spinner = TTY::Spinner.new( ':spinner :title', clear: true )
        spinner.update( title: "Collecting camera info.." )
        spinner.auto_spin

        # run in a thread so the spinner spins
        Thread.new {
            cameras = GPhoto2::Camera.all

            # wait for a camera, if requested
            if self.config[ :wait_for_camera ] && cameras.empty?
                spinner.update( title: "Waiting for a camera.." )

                loop do
                    cameras = GPhoto2::Camera.all
                    break unless cameras.empty?
                    sleep 0.5
                end
            end

            begin
                case cameras.length
                when 0
                    raise NoCameraError
                when 1
                    begin
                        self.selected_camera = cameras.first
                        self.selected_camera.reload
                        spinner.stop
                    rescue
                        # not a valid camera - reload throws unsupported error
                        raise NoCameraError
                    end
                else
                    choices = cameras.each_with_object( {} ) do |c,acc|
                            deets = c rescue nil
                            model = deets['cameramodel'] rescue nil
                            next unless model && deets
                            acc.merge!({ model => deets })
                        end

                    spinner.stop

                    # if one (or more) of the choices were invalid, and left us with only one
                    # choice, just choose that one. this filters out "USB Mass Storage Camera"
                    if choices.length == 1
                        self.selected_camera = choices.first.last
                        self.selected_camera.reload
                    else
                        choices.merge!( { "Exit" => "exit" } )

                        self.selected_camera = self.prompt.select(
                            "Which camera?",
                            choices
                        )
                        exit 0 if self.selected_camera == "exit"
                    end
                end
            rescue NoCameraError
                spinner.stop
                puts "No cameras found; did you turn it on?"
                exit 2
            end
        }.join
    end


    ##############################################
    # T H R E A D  F U N C T I O N S
    ##############################################

    # run the middleman thread that makes sure our capture keeps going, etc
    # this will start the capture thread for us.
    #
    def run_interface_thread
        self.interface_thread = Thread.new do
            @timing_out = nil
            loop do
                unless self.capture_thread && self.capture_thread.alive?
                    # if the thread died with a message on the bus, make sure
                    # it's not meant to end the program
                    item = self.result( fail_ok: true )
                    if item && item.first == "abort"
                        print TTY::Cursor.clear_lines(10, :down)
                        self.prompt.say "Abort: camera disconnected? (%s)" % [ item.last ]
                        exit 3
                    end

                    # otherwise, just (re)start it.
                    self.capture_thread = Thread.new{ self.run_capture_thread }
                end

                # close the capture on timeout (nobody is streaming)
                pids = IO.popen( ['lsof', '-w', '-Fp', self.config[ :output_device ]], "r", :err=>[:child, :out] ).readlines.delete_if{|x| x=='' } rescue []
                if pids.count <= 1
                    @timing_out ||= Time.now()
                    if ( Time.now().to_i - @timing_out.to_i ) > self.config[ :capture_timeout ]
                        print TTY::Cursor.clear_lines(10, :down)
                        self.prompt.say 'Capture Timeout: No clients using stream for %s minutes.' % [ self.config[ :capture_timeout ] / 60 ]
                        exit 4
                    end
                else
                    @timing_out = nil
                end

                sleep 0.5
            end
        end
    end


    # the meat'n'potatoes; this will be started in a new thread automatically
    # spawns ffmpeg, starts a capture, and listens on $in_queue for commands
    #
    def run_capture_thread
        # capture a trial frame so we can get details about the camera's output
        image = Magick::Image.from_blob( self.selected_camera.preview.data ).first
        self.capture_geometry_width  = image.columns
        self.capture_geometry_height = image.rows
        self.capture_geometry_aspect = Rational( image.columns, image.rows )

        io = IO.popen( self.ffmpeg_cmdline, "w" ) do |ffmpeg_io|
            loop do
                # read a preview image and send it off to ffmpeg
                ffmpeg_io.write( self.selected_camera.preview.data )

                # attempt to load/run a new task
                action, data = begin
                    self.command_queue.pop( true )
                rescue
                    [ nil, nil ]
                end

                case action
                    when 'autofocusdrive'
                        self.selected_camera.update( autofocusdrive: 1 )
                        #self.selected_camera.update( eosremoterelease: "Press Half" )
                        #self.selected_camera.update( eosremoterelease: "Release Full" )
                        #self.selected_camera.update( autofocusdrive: 0 )
                        #self.selected_camera.update( eosremoterelease: "None" )

                    when 'cancelautofocus'
                        # i honestly have no idea what it's doing, or what i'm doing
                        self.selected_camera.update( autofocusdrive: 0 )
                        self.selected_camera.update( cancelautofocus: 1 )
                        self.selected_camera.update( cancelautofocus: 0 )

                    when 'manualfocusdrive'
                        self.selected_camera.update( manualfocusdrive: data )
                        self.selected_camera.update( manualfocusdrive: "None" )

                    when 'getconfig'
                        node,func = data.split( ":" )
                        if node && func
                            self.result_queue.push( [ self.selected_camera.config[ node ].send( func.to_sym ) ] )
                        else
                            self.result_queue.push( [ self.selected_camera.config ] )
                        end

                    when 'update'
                        node,value = data.split(":")
                        self.selected_camera.update( node.to_sym => value ) rescue nil

                    when 'exit'
                        self.result_queue.push [ 'bye' ]
                        break
                end
            end # loop
        end # io

    rescue Errno::EPIPE
        # ffmpeg crash
        self.result_queue.push( [ 'restart_pipe' ] ) rescue nil

    rescue GPhoto2::Error => e
        self.result_queue.push( [ 'abort', e.message ] ) rescue nil

    ensure
        # signal a close to gphoto and kill ffmpeg
        self.selected_camera.exit

        pstatus = Process.kill(0, io.pid ) rescue nil
        if pstatus
            Process.kill( 9, io.pid )
            Process.waitpid io.pid
        end
    end


    ##############################################
    # V I D E O  F U N C T I O N S
    ##############################################

    # make sure the v4l2 device is loaded
    #
    def check_v4l2
        unless File.exist?( self.config[ :output_device ] )
            puts "No such file: %s - is v4l2loopback loaded?"
            exit 1
        end
    end


    # The command used to run ffmpeg, which ingests the preview frames from the camera.
    # TODO: flesh out the ffmpeg customization.
    #
    def ffmpeg_cmdline
        return "ffmpeg -hide_banner -loglevel quiet -i - -vcodec rawvideo -pix_fmt yuv420p -vf scale=w=%s:h=%s%s -threads 0 -f v4l2 %s" % [
                self.config[ :output_geometry_width ],
                self.config[ :output_geometry_height ],
                ( self.config[ :output_geometry_keep_aspect ] && self.config[ :output_geometry_aspect ] != self.capture_geometry_aspect ?
                    ( self.config[ :output_geometry_aspect ] > self.capture_geometry_aspect ? ':force_original_aspect_ratio=increase' : ':force_original_aspect_ratio=decrease' ) : ''
                ),
                self.config[ :output_device ]
            ]
    end


    ##############################################
    # D A T A  Q U E U E  F U N C T I O N S
    ##############################################

    # send a command to the processing thread, optionally
    # waiting for and returning a result inline
    #
    def command( c, arg: nil, wait_for_result: false, fail_ok: false )
        cr = self.command_queue.push( [ c, arg ].flatten )

        return self.result( fail_ok: fail_ok ) if wait_for_result
        return cr
    end


    # get a result from the queue, optionally allowing for failures
    # (i.e. empty queue)
    #
    def result( fail_ok:false )
        begin
            return self.result_queue.pop( fail_ok )
        rescue => e
            # re-raise if fail_ok=>false
            raise e unless fail_ok
            return nil
        end
    end


    ##############################################
    # M E N U  F U N C T I O N S
    ##############################################

    # run the main menu in an infinite loop, until exited.
    #
    def menu_main
        main_menu = [
            { name: "Autofocus", key: 'f',value: -> { self.command( 'autofocusdrive' ); 1 } },
            { name: "Cancel Autofocus", key: 'c', value: -> { self.command( 'cancelautofocus' ); 2 } },
            { name: "Manual Focus", key: 'm', value: -> { self.menu_manualfocus; 3 } },
            { name: "Adjust Aperture", key: 'a', value: -> {
                    self.menu_config_slider( 'aperture', 'an aperture value.' ) do |values|
                        pretty = values.map{|v| "f/%s" % [ v ] }

                        [ values, pretty ]
                    end
                    4
                } },
            { name: "Adjust Shutter Speed", key: 's', value: -> { self.menu_config_slider( 'shutterspeed', 'a shutter speed.' ); 5 } },
            { name: "Adjust ISO", key: 'i', value: -> { self.menu_config_slider( 'iso', 'an ISO.' ); 6 } },
            { name: "Send gphoto2 config", key: 'g', value: Proc.new{
                    cmd = self.prompt.ask "Enter a gphoto2 config like \"node_name:value\". Enter an empty line to cancel.\n>"
                    self.command( 'update', arg: cmd ) if cmd
                    7
                }
            },
            { name: "Restart capture", key: "r", value: -> { self.command( 'exit', wait_for_result: true ); 8} },
            { name: "Exit", key: :escape, key_name: 'esc', value: -> { self.interface_thread.kill;self.interface_thread.join;self.capture_thread.kill;exit 0 } }
          ]

        selected = 1
        loop do
            selected = self.prompt.select( "Camera Capture", main_menu, per_page: 10, default: selected, enum: '.', key_action: :select )
        end
    end


    # automatically create a 'select' menu for a config (shutterspeed, apeture, etc..)
    #
    def menu_config_select( configname, configtitle, &block )
        # get valid & current values
        values = self.command( 'getconfig', arg: '%s:choices' % [configname], wait_for_result: true ).first
        pretty_values = values

        # if a block was given, use it to sort/translate/etc the values
        # otherwise, just use the values as-is
        values, pretty_values = block.call( values ) if block_given?

        choices = values.each_with_object( [] ).with_index do |(val, acc), idx|
            acc << { name: pretty_values[ idx ], value: val }
        end

        # add the "exit" option
        choices << { name: "Back to main menu", value: "Back to main menu" }

        loop do
            current = self.command( 'getconfig', arg: '%s:value' % [configname], wait_for_result: true ).first
            default = choices.find_index{|c| c[ :value ] == current } + 1

            selected = self.prompt.select( "Select %s." % [configtitle], choices, default: default, per_page: 15, filter: true )
            break if selected =~ /^Back/

            self.command( 'update', arg: '%s:%s' % [ configname, selected ] )
        end
    end


    # automatically create a 'slider' menu for a config (shutterspeed, apeture, etc..)
    # note: this isn't currently working, and is pending the next release (>0.22) of tty-prompt
    #
    def menu_config_slider( configname, configtitle, &block )
        # get valid & current values
        values = self.command( 'getconfig', arg: '%s:choices' % [configname], wait_for_result: true ).first
        pretty_values = values

        # if a block was given, use it to sort/translate/etc the values
        # otherwise, just use the values as-is
        values, pretty_values = block.call( values ) if block_given?

        choices = values.each_with_object( [] ).with_index do |(val, acc), idx|
            acc << { name: pretty_values[ idx ], value: val }
        end

        # add the "exit" option
        choices << { name: "Exit", value: "Exit", key: :escape }

        loop do
            current = self.command( 'getconfig', arg: '%s:value' % [configname], wait_for_result: true ).first
            default = choices.find_index{|c| c[ :value ] == current } + 1

            selected = self.prompt.slider( "Select %s." % [configtitle], choices, default: default, show_help: 'always' )
            break if selected == "Exit"

            self.command( 'update', arg: '%s:%s' % [ configname, selected ] )
        end
    end


    # the menu for adjusting the focus manually
    #
    def menu_manualfocus
        last_selection = "None"

        # cancel auto focus first to regain control
        af_status = self.command( 'getconfig', arg: 'autofocusdrive:value', wait_for_result: true ).first
        if af_status
            self.command( 'cancelautofocus' )
        end

        # grab the focus types
        focus_modes = self.command( 'getconfig', arg: 'manualfocusdrive:choices', wait_for_result: true ).first

        # we want to sort near 3..1, none, far 1..3
        focus_modes.sort_by!{|n|
            weight = 0
            if n != "None"
                type,amt = n.split(' ')
                weight = weight + ( type == "Near" ? amt.to_i*-1 : amt.to_i )
            end

            weight
        }

        # then sub the numbers with plus signs
        focus_modes = focus_modes.map{|n|
           type,amt = n.split(' ')
            if amt
                {name: "%s %s" % [ type, "+"*amt.to_i ], value: n}
            else
                {name: n, value: n}
            end
        }

        focus_modes << "Back to main menu"

        loop do
            selected = self.prompt.select(
                "Select a focus adjustment",
                focus_modes,
                per_page: 10,
                default: focus_modes.find_index{|n| n[:value] == last_selection }+1
            )
            break if selected =~ /^Back/
            self.command( 'manualfocusdrive', arg: selected )
            last_selection = selected
        end
    end

end


# gather commandline options
options = {}
OptionParser.new do |opts|
    opts.banner = "Usage: #{ARGV[0]} [options]"

    opts.on( '-cCONFIG', '--config=CONFIG', "The location of the camera-capture config file. (default: ~/.config/camera-capture.yml)") do |c|
        options[ :config ] = Pathname( c )
    end

    opts.on( '-dDEVICE', '--device=DEVICE', "The v4l2 device to output to. (default: video9)") do |d|
        options[ :output_device ] = d
    end

    opts.on( '-g', '--geometry=WxH', "Specify the output video geometry. (default: 1920x1280)") do |g|
        matches = /^(?<w>\d+)[xX:](?<h>\d+)$/.match( g )
        raise ArgumentError, "Invalid geometry: %p" % [ g ] unless matches

        options[ :output_geometry_w ] = matches[ :w ]
        options[ :output_geometry_h ] = matches[ :h ]
    end

    opts.on( '-tMINUTES', '--timeout=minutes', "Automatically stop the capture if no client(s) are consuming the stream after MINUTES. (default: 5)") do |m|
        matches = /^\d+$/.match(m)
        raise ArgumentError, "Invalid timeout: %p" % [ m ] unless matches

        options[ :capture_timeout ] = ( m.to_i * 60 )
    end

    opts.on( '-w', '--[no-]wait', "Wait for a camera to be connected. (default: no-wait)") do |w|
        options[ :wait_for_camera ] = w
    end

    opts.on( '--[no-]keep-aspect-ratio', "Perserve the aspect ratio of the captured video in the output stream. (default: true)" ) do |k|
        options[ :output_geometry_keep_aspect ] = k
    end

    opts.on( '-h', '--help', "Print this help.") do |h|
        puts opts
        exit 0
    end
end.parse!


# run!
app = CameraCapture.new( **options )
app.run