gallery-thumbgen.lua

--[[
mpv-gallery-view | https://github.com/occivink/mpv-gallery-view

This mpv script implements a worker for generating gallery thumbnails.
It is meant to be used by other scripts.
Multiple copies of this script can be loaded by mpv.

File placement: inside scripts directory
Settings: script-opts/gallery_worker.conf
]]

local utils = require 'mp.utils'
local msg = require 'mp.msg'

local jobs_queue = {} -- queue of thumbnail jobs
local failed = {} -- list of failed output paths, to avoid redoing them
local script_id = mp.get_script_name() .. utils.getpid()

local opts = {
    ytdl_exclude = "",
}
(require 'mp.options').read_options(opts, "gallery_worker")

local ytdl = {
    path = "youtube-dl",
    searched = false,
    blacklisted = {} -- Add patterns of URLs you want blacklisted from youtube-dl,
                     -- see gallery_worker.conf or ytdl_hook-exclude in the mpv manpage for more info
}

function append_table(lhs, rhs)
    for i = 1,#rhs do
        lhs[#lhs+1] = rhs[i]
    end
    return lhs
end

local function file_exists(path)
    local info = utils.file_info(path)
    return info ~= nil and info.is_file
end

local video_extensions = { "mkv", "webm", "mp4", "avi", "wmv" }

function is_video(input_path)
    local extension = string.match(input_path, "%.([^.]+)$")
    if extension then
        extension = string.lower(extension)
        for _, ext in ipairs(video_extensions) do
            if extension == ext then
                return true
            end
        end
    end
    return false
end

function is_blacklisted(url)
    if opts.ytdl_exclude == "" then return false end
    if #ytdl.blacklisted == 0 then
        local joined = opts.ytdl_exclude
        while joined:match('%|?[^|]+') do
            local _, e, substring = joined:find('%|?([^|]+)')
            table.insert(ytdl.blacklisted, substring)
            joined = joined:sub(e+1)
        end
    end
    if #ytdl.blacklisted > 0 then
        url = url:match('https?://(.+)')
        for _, exclude in ipairs(ytdl.blacklisted) do
            if url:match(exclude) then
                msg.verbose('URL matches excluded substring. Skipping.')
                return true
            end
        end
    end
    return false
end


function ytdl_thumbnail_url(input_path)
    local function exec(args)
        local ret = utils.subprocess({args = args, cancellable=false})
        return ret.status, ret.stdout, ret
    end
    local function first_non_nil(x, ...)
        if x ~= nil then return x end
        return first_non_nil(...)
    end

    -- if input_path is youtube, generate our own URL
    youtube_id1 = string.match(input_path, "https?://youtu%.be/([%a%d%-_]+).*")
    youtube_id2 = string.match(input_path, "https?://w?w?w?%.?youtube%.com/v/([%a%d%-_]+).*")
    youtube_id3 = string.match(input_path, "https?://w?w?w?%.?youtube%.com/watch%?v=([%a%d%-_]+).*")
    youtube_id4 = string.match(input_path, "https?://w?w?w?%.?youtube%.com/embed/([%a%d%-_]+).*")
    youtube_id = youtube_id1 or youtube_id2 or youtube_id3 or youtube_id4

    if youtube_id then
        -- the hqdefault.jpg thumbnail should always exist, since it's used on the search result page
        return "https://i.ytimg.com/vi/" .. youtube_id ..  "/hqdefault.jpg"
    end

    --otherwise proceed with the slower `youtube-dl -J` method
    if not (ytdl.searched) then --search for youtude-dl in mpv's config directory
        local exesuf = (package.config:sub(1,1) == '\\') and '.exe' or ''
        local ytdl_mcd = mp.find_config_file("youtube-dl")
        if not (ytdl_mcd == nil) then
            msg.error("found youtube-dl at: " .. ytdl_mcd)
            ytdl.path = ytdl_mcd
        end
        ytdl.searched = true
    end
    local command = {ytdl.path, "--no-warnings", "--no-playlist", "-J", input_path}
    local es, json, result = exec(command)

    if (es < 0) or (json == nil) or (json == "") then
        msg.error("fetching thumbnail url with youtube-dl failed for" .. input_path)
        return input_path
    end
    local json, err = utils.parse_json(json)
    if (json == nil) then
        msg.error("failed to parse json for youtube-dl thumbnail: " .. err)
        return input_path
    end

    if (json.thumbnail == nil) or (json.thumbnail == "") then
        msg.error("no thumbnail url from youtube-dl.")
        return input_path
    end
    return json.thumbnail
end

function thumbnail_command(input_path, width, height, take_thumbnail_at, output_path, accurate, with_mpv)
    local vf = string.format("%s,%s",
        string.format("scale=iw*min(1\\,min(%d/iw\\,%d/ih)):-2", width, height),
        string.format("pad=%d:%d:(%d-iw)/2:(%d-ih)/2:color=0x00000000", width, height, width, height)
    )
    local out = {}
    local add = function(table) out = append_table(out, table) end


    if input_path:find("^https?://") and not is_blacklisted(input_path) then
        -- returns the original input_path on failure
        input_path = ytdl_thumbnail_url(input_path)
    end

    if input_path:find("^archive://") or input_path:find("^edl://") then
        with_mpv = true
    end


    if not with_mpv then
        out = { "ffmpeg" }
        if is_video(input_path) then
            if string.sub(take_thumbnail_at, -1) == "%" then
                --if only fucking ffmpeg supported percent-style seeking
                local res = utils.subprocess({ args = {
                    "ffprobe", "-v", "error",
                    "-show_entries", "format=duration", "-of",
                    "default=noprint_wrappers=1:nokey=1", input_path
                }, cancellable = false })
                if res.status == 0 then
                    local duration = tonumber(string.match(res.stdout, "^%s*(.-)%s*$"))
                    if duration then
                        local percent = tonumber(string.sub(take_thumbnail_at, 1, -2))
                        local start = tostring(duration * percent / 100)
                        add({ "-ss", start })
                    end
                end
            else
                add({ "-ss", take_thumbnail_at })
            end
        end
        if not accurate then
            add({"-noaccurate_seek"})
        end
        add({
            "-i", input_path,
            "-vf", vf,
            "-map", "v:0",
            "-f", "rawvideo",
            "-pix_fmt", "bgra",
            "-c:v", "rawvideo",
            "-frames:v", "1",
            "-y", "-loglevel", "quiet",
            output_path
        })
    else
        out = { "mpv", input_path }
        if take_thumbnail_at ~= "0" and is_video(input_path) then
            if not accurate then
                add({ "--hr-seek=no"})
            end
            add({ "--start="..take_thumbnail_at })
        end
        add({
            "--no-config", "--msg-level=all=no",
            "--vf=lavfi=[" .. vf .. ",format=bgra]",
            "--audio=no",
            "--sub=no",
            "--frames=1",
            "--image-display-duration=0",
            "--of=rawvideo", "--ovc=rawvideo",
            "--o="..output_path
        })
    end
    return out
end

function generate_thumbnail(thumbnail_job)
    if file_exists(thumbnail_job.output_path) then return true end

    local dir, _ = utils.split_path(thumbnail_job.output_path)
    local tmp_output_path = utils.join_path(dir, script_id)

    local command = thumbnail_command(
        thumbnail_job.input_path,
        thumbnail_job.width,
        thumbnail_job.height,
        thumbnail_job.take_thumbnail_at,
        tmp_output_path,
        thumbnail_job.accurate,
        thumbnail_job.with_mpv
    )

    local res = utils.subprocess({ args = command, cancellable = false })
    --"atomically" generate the output to avoid loading half-generated thumbnails (results in crashes)
    if res.status == 0 then
        local info = utils.file_info(tmp_output_path)
        if not info or not info.is_file or info.size == 0 then
            return false
        end
        if os.rename(tmp_output_path, thumbnail_job.output_path) then
            return true
        end
    end
    return false
end

function handle_events(wait)
    e = mp.wait_event(wait)
    while e.event ~= "none" do
        if e.event == "shutdown" then
            return false
        elseif e.event == "client-message" then
            if e.args[1] == "push-thumbnail-front" or e.args[1] == "push-thumbnail-back" then
                local thumbnail_job = {
                    requester = e.args[2],
                    input_path = e.args[3],
                    width = tonumber(e.args[4]),
                    height = tonumber(e.args[5]),
                    take_thumbnail_at = e.args[6],
                    output_path = e.args[7],
                    accurate = (e.args[8] == "true"),
                    with_mpv = (e.args[9] == "true"),
                }
                if e.args[1] == "push-thumbnail-front" then
                    jobs_queue[#jobs_queue + 1] = thumbnail_job
                else
                    table.insert(jobs_queue, 1, thumbnail_job)
                end
            end
        end
        e = mp.wait_event(0)
    end
    return true
end

local registration_timeout = 2 -- seconds
local registration_period = 0.2

-- shitty custom event loop because I can't figure out a better way
-- works pretty well though
function mp_event_loop()
    local start_time = mp.get_time()
    local sleep_time = registration_period
    local last_broadcast_time = -registration_period
    local broadcast_func
    broadcast_func = function()
        local now = mp.get_time()
        if now >= start_time + registration_timeout then
            mp.commandv("script-message", "thumbnails-generator-broadcast", mp.get_script_name())
            sleep_time = 1e20
            broadcast_func = function() end
        elseif now >= last_broadcast_time + registration_period then
            mp.commandv("script-message", "thumbnails-generator-broadcast", mp.get_script_name())
            last_broadcast_time = now
        end
    end

    while true do
        if not handle_events(sleep_time) then return end
        broadcast_func()
        while #jobs_queue > 0 do
            local thumbnail_job = jobs_queue[#jobs_queue]
            if not failed[thumbnail_job.output_path] then
                if generate_thumbnail(thumbnail_job) then
                    mp.commandv("script-message-to", thumbnail_job.requester, "thumbnail-generated", thumbnail_job.output_path)
                else
                    failed[thumbnail_job.output_path] = true
                end
            end
            jobs_queue[#jobs_queue] = nil
            if not handle_events(0) then return end
            broadcast_func()
        end
    end
end