Add auto transcoding
This commit is contained in:
parent
ae855fc1c5
commit
7fedb753c6
GPG key ID:
0FA334385D0B689F
3 changed files with 208 additions and 3 deletions
|
|
@ -7,15 +7,20 @@ log = logging.getLogger(__name__)
|
|||
|
||||
|
||||
@click.group()
|
||||
def cli():
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
@click.option(
|
||||
"--log-level",
|
||||
default="INFO",
|
||||
type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False),
|
||||
help="Set the log level."
|
||||
)
|
||||
def cli(log_level: str):
|
||||
logging.basicConfig(level=log_level.upper())
|
||||
|
||||
|
||||
@cli.command()
|
||||
def run():
|
||||
"""Runs the bot"""
|
||||
log.info("Starting bot.")
|
||||
logging.getLogger("nio.rooms").setLevel(logging.WARNING)
|
||||
from .main import bot
|
||||
|
||||
run_async(bot.start(access_token=bot.cfg["bot"]["access_token"]))
|
||||
|
|
|
|||
|
|
@ -8,6 +8,10 @@ import tortoise
|
|||
from pathlib import Path
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
logging.getLogger("nio.rooms").setLevel(logging.WARNING)
|
||||
logging.getLogger("nio.crypto.log").setLevel(logging.WARNING)
|
||||
logging.getLogger("peewee").setLevel(logging.WARNING)
|
||||
logging.getLogger("nio.responses").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
with open("config.toml", "rb") as fd:
|
||||
|
|
@ -73,3 +77,7 @@ async def on_command_error(ctx: niobot.Context, exc: Exception):
|
|||
@bot.on_event("command_complete")
|
||||
async def on_command_complete(ctx: niobot.Context, _):
|
||||
log.info("Command %s completed.", ctx.command)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tortoise.run_async(bot.start(access_token=bot.cfg["bot"]["access_token"]))
|
||||
|
|
|
|||
192
app/modules/ts_transcode.py
Normal file
192
app/modules/ts_transcode.py
Normal file
|
|
@ -0,0 +1,192 @@
|
|||
import logging
|
||||
import shutil
|
||||
import typing
|
||||
import urllib.parse
|
||||
|
||||
import niobot
|
||||
import tempfile
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class TruthSocialTranscode(niobot.Module):
|
||||
USER_IDS = {
|
||||
"@trumpbot:shronk.net",
|
||||
"@tatebot:shronk.net",
|
||||
"@nex:nexy7574.co.uk"
|
||||
}
|
||||
|
||||
def __init__(self, bot: niobot.NioBot):
|
||||
super().__init__(bot)
|
||||
self.log = logging.getLogger(__name__)
|
||||
if not shutil.which("ffprobe"):
|
||||
raise RuntimeError("ffprobe is required for this module.")
|
||||
if not shutil.which("ffmpeg"):
|
||||
raise RuntimeError("ffmpeg is required for this module.")
|
||||
|
||||
# noinspection PyTypeChecker
|
||||
self.bot.add_event_callback(self.on_message, niobot.RoomMessageVideo)
|
||||
|
||||
def __teardown__(self):
|
||||
for callback in self.bot.event_callbacks:
|
||||
if callback.func == self.on_message:
|
||||
self.bot.event_callbacks.remove(callback)
|
||||
super().__teardown__()
|
||||
|
||||
async def on_message_internal(self, room: niobot.MatrixRoom, event: niobot.RoomMessageVideo) -> str:
|
||||
self.log.info("Processing event: %r (%r)", event, room)
|
||||
# if not isinstance(event, niobot.RoomMessageVideo):
|
||||
# # We are not interested in non-videos.
|
||||
# return
|
||||
# event: niobot.RoomMessageVideo
|
||||
|
||||
if event.sender not in self.USER_IDS:
|
||||
# We are not interested in messages from other users.
|
||||
self.log.info("Ignored user: %s", event.sender)
|
||||
return "IGNORED"
|
||||
|
||||
# Videos sent in the truth room are not encrypted, so that makes this easier.
|
||||
http_url = await self.bot.mxc_to_http(
|
||||
event.url,
|
||||
homeserver=event.sender.split(":", 1)[1] # use the sender homeserver for download
|
||||
)
|
||||
http_url = "https://" + http_url
|
||||
name = event.flattened()["content.filename"]
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix="-%s" % name) as temp_fd:
|
||||
# async with self.bot.client_session.get(http_url) as resp:
|
||||
# temp_fd.write(await resp.read())
|
||||
# temp_fd.flush()
|
||||
dl = await self.bot.download(event.url, save_to=Path(temp_fd.name))
|
||||
if isinstance(dl, niobot.DownloadError):
|
||||
self.log.error("Error downloading video: %r", dl)
|
||||
await self.bot.send_message(
|
||||
room,
|
||||
"Error downloading video, unable to transcode: %r" % dl,
|
||||
reply_to=event.event_id
|
||||
)
|
||||
return "ERROR_DOWNLOAD"
|
||||
|
||||
# First, probe the video to see if it has H265 content in it.
|
||||
probe = await asyncio.create_subprocess_exec(
|
||||
"ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
|
||||
"stream=codec_name", "-of", "default=noprint_wrappers=1:nokey=1",
|
||||
temp_fd.name,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
stdout, stderr = await probe.communicate()
|
||||
if probe.returncode != 0:
|
||||
# Something went wrong, let's not continue.
|
||||
self.log.error("Error probing video:\n%s", (stderr or b"EOF").decode("utf-8"))
|
||||
await self.bot.send_message(
|
||||
room,
|
||||
"Error probing video, unable to determine whether a transcode is required.",
|
||||
reply_to=event.event_id
|
||||
)
|
||||
return "ERROR_PROBE"
|
||||
|
||||
codec = stdout.decode("utf-8").strip()
|
||||
if codec in ["h264", "vp8", "vp9"]:
|
||||
# We don't need to transcode this.
|
||||
self.log.info("Video is already in a supported format: %s", codec)
|
||||
return "NO_NEED"
|
||||
|
||||
# We now need to transcode this. For compatibility, we'll force it to be H264 and AAC.
|
||||
# We will, however, use the High profile for H264, as almost everything supports it now, plus better
|
||||
# compression.
|
||||
# We will also limit the bitrate to about 5 megabit, because Trump has a habit of posting raw footage
|
||||
# that ends up being something stupid like 20 megabit at 60fps.
|
||||
with tempfile.NamedTemporaryFile(suffix="-transcoded-" + name + ".mp4") as out_fd:
|
||||
self.log.info(
|
||||
"transcoding %s -> %s",
|
||||
temp_fd.name,
|
||||
out_fd.name
|
||||
)
|
||||
args = [
|
||||
"ffmpeg",
|
||||
"-hwaccel",
|
||||
"auto",
|
||||
"-v",
|
||||
"error",
|
||||
"-i",
|
||||
temp_fd.name,
|
||||
"-c:v",
|
||||
"libx264",
|
||||
"-pix_fmt",
|
||||
"yuv420p",
|
||||
"-profile:v",
|
||||
"high",
|
||||
"-preset",
|
||||
"medium",
|
||||
"-b:v",
|
||||
"5M",
|
||||
"-maxrate",
|
||||
"5M",
|
||||
"-bufsize",
|
||||
"10M",
|
||||
"-minrate",
|
||||
"1M",
|
||||
"-c:a",
|
||||
"aac",
|
||||
"-b:a",
|
||||
"96k",
|
||||
"-vf",
|
||||
"scale=-2:'min(1080,ih)'",
|
||||
"-movflags",
|
||||
"faststart",
|
||||
"-y",
|
||||
"-f",
|
||||
"mp4",
|
||||
out_fd.name
|
||||
]
|
||||
transcode = await asyncio.create_subprocess_exec(
|
||||
*args,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
self.log.info("Starting transcode...")
|
||||
stdout, stderr = await transcode.communicate()
|
||||
self.log.info("Finished transcode.")
|
||||
if transcode.returncode != 0:
|
||||
self.log.error("Error transcoding video:\n%s", stderr.decode("utf-8"))
|
||||
await self.bot.send_message(
|
||||
room,
|
||||
"Error transcoding video, unable to transcode: got return code %d" % transcode.returncode,
|
||||
reply_to=event.event_id
|
||||
)
|
||||
return "ERROR_TRANSCODE"
|
||||
# We transcoded! Now we need to upload it.
|
||||
file = await niobot.VideoAttachment.from_file(out_fd.name)
|
||||
self.log.info("Uploading transcoded video")
|
||||
await self.bot.send_message(
|
||||
room.room_id,
|
||||
file=file,
|
||||
reply_to=event.event_id
|
||||
)
|
||||
self.log.info("Uploaded transcoded video")
|
||||
return "OK"
|
||||
|
||||
async def on_message(self, room, event):
|
||||
try:
|
||||
return await self.on_message_internal(room, event)
|
||||
except Exception as e:
|
||||
self.log.exception("Error processing video", exc_info=e)
|
||||
return "ERROR %r" % e
|
||||
|
||||
@staticmethod
|
||||
def is_allowed(ctx: niobot.Context) -> bool:
|
||||
return ctx.message.sender.split(":", 1)[1] in {"shronk.net", "nicroxio.co.uk", "nexy7574.co.uk"}
|
||||
|
||||
@niobot.command("transcode")
|
||||
async def do_transcode(self, ctx: niobot.Context, message_link: typing.Annotated[niobot.RoomMessage, niobot.EventParser("m.room.message")]):
|
||||
"""
|
||||
Transcodes a video to H264/AAC.
|
||||
"""
|
||||
if not isinstance(message_link, niobot.RoomMessageVideo):
|
||||
await ctx.respond("That's not a video. %r" % message_link)
|
||||
return
|
||||
res = await ctx.respond("\N{hourglass} Dispatching, please wait...")
|
||||
ret = await self.on_message(ctx.room, message_link)
|
||||
await res.edit(f"\N{white heavy check mark} Transcode complete. Result: `{ret}`")
|
||||
|
||||
Loading…
Reference in a new issue