Advertisement
CustomUrlsSuck

Untitled

Oct 2nd, 2022
902
0
Never
1
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 12.00 KB | None | 0 0
  1. #!/usr/bin/env python3
  2.  
  3. import time
  4. import ffmpeg
  5. import os
  6. import sys
  7. import docopt
  8. import subprocess
  9. import datetime
  10. import re
  11. import asyncio
  12. import urwid
  13. from collections import deque
  14.  
  15. USAGE = """Reencode Videos
  16. Usage:
  17.    ./fbed.py <parallel_encodes> <items>...
  18.  
  19. Guide:
  20.    <items> can be a single files or a directories. If a directory is passed all
  21.    files in the directory and it subdirectories besides those in one named
  22.    'encode_output' will be re-encoded
  23. """
  24.  
  25. output_dir = "encode_output"
  26. match_out_time = re.compile("(\d+):(\d+):(\d+\.\d+)")
  27.  
  28. def parse_out_time(out_time):
  29.     m = match_out_time.match(out_time)
  30.     hours = int(m.group(1))
  31.     minutes = int(m.group(2))
  32.     seconds_millis = float(m.group(3))
  33.     seconds = int(seconds_millis)
  34.     milliseconds = int((seconds_millis - seconds) * 1000)
  35.     return datetime.timedelta(hours=hours, minutes=minutes, seconds=seconds, milliseconds=milliseconds)
  36.  
  37. def get_video_bitrate(probe):
  38.     info = [s for s in probe["streams"] if s["codec_type"] == "video"][0]
  39.     if "bit_rate" in info:
  40.         return int(int(info["bit_rate"]) / 1000)
  41.     elif "bit_rate" in probe["format"]:
  42.         return int(int(probe["format"]["bit_rate"]) / 1000)
  43.     else:
  44.         print("Failed to read bitrate from ffprobe! Please include the information below in your Github issue")
  45.         print(probe)
  46.         sys.exit(1)
  47.  
  48. class EncodingTask:
  49.     def __init__(self, filename, out_filename):
  50.         os.makedirs(os.path.dirname(out_filename), exist_ok=True)
  51.         self.out_filename = out_filename
  52.         self.log_filename = os.path.splitext(self.out_filename)[0] + ".log"
  53.         self.stderr = open(self.log_filename, "w", encoding="utf8")
  54.         self.pipe_read, self.pipe_write = os.pipe()
  55.         self.pipe_read_file = os.fdopen(self.pipe_read)
  56.  
  57.         probe = ffmpeg.probe(filename)
  58.         duration = float(probe["format"]["duration"])
  59.         seconds = int(duration)
  60.         milliseconds = int((duration - seconds) * 1000)
  61.         self.duration = datetime.timedelta(seconds=seconds, milliseconds=milliseconds)
  62.  
  63.         info = [s for s in probe["streams"] if s["codec_type"] == "video"][0]
  64.         self.width = info["width"]
  65.         self.height = info["height"]
  66.         source_bitrate = get_video_bitrate(probe)
  67.         # Pick bitrate based on resolution, 1080p (8Mbps), 720p (5Mbps), smaller (3Mbps)
  68.         bitrate = 3000
  69.         if self.height > 720:
  70.             bitrate = 8000
  71.         elif self.height > 480:
  72.             bitrate = 5000
  73.         # Don't exceed the source bitrate as our target
  74.         if bitrate > source_bitrate:
  75.             bitrate = source_bitrate
  76.  
  77.         encoding_args = {
  78.             # HWAccel for RPi4, may need to pick a different encoder
  79.             # for HW accel on other systems
  80.             "c:v": "h264_v4l2m2m",
  81.             "num_output_buffers": 32,
  82.             "num_capture_buffers": 16,
  83.             "b:v": f"{bitrate}k",
  84.             "c:a": "copy",
  85.             "progress": f"pipe:{self.pipe_write}"
  86.         }
  87.         self.start = datetime.datetime.now()
  88.         in_stream = ffmpeg.input(filename)
  89.         video = in_stream.video.filter("format", **{"pix_fmts": "yuv420p"})
  90.         enc = ffmpeg.output(video, in_stream.audio, self.out_filename, **encoding_args)
  91.         args = ffmpeg.compile(enc, overwrite_output=True)
  92.         self.proc = subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=self.stderr, pass_fds=[self.pipe_write])
  93.         self.encode_stats = {}
  94.         self.encode_error = False
  95.  
  96.     def is_complete(self):
  97.         encode_done = False
  98.         while not encode_done and not self.encode_error:
  99.             # TODO: Needs some non-blocking reading here in case the process dies while
  100.             # we're trying to read progress
  101.             l = self.pipe_read_file.readline()
  102.             if not l:
  103.                 break
  104.  
  105.             l = l.strip()
  106.             key, val = l.split("=")
  107.             val = val.strip()
  108.             if key == "out_time":
  109.                 out_time = parse_out_time(val)
  110.                 self.encode_stats["percent_done"] = (100.0 * out_time.total_seconds()) / self.duration.total_seconds()
  111.                 self.encode_stats[key] = out_time
  112.             elif key == "fps":
  113.                 self.encode_stats[key] = float(val)
  114.             else:
  115.                 self.encode_stats[key] = val
  116.  
  117.             if l.startswith("progress="):
  118.                 encode_done = l == "progress=end"
  119.                 break
  120.  
  121.         if not encode_done:
  122.             speed = float(self.encode_stats["speed"][:-1])
  123.             remaining_time = (self.duration - self.encode_stats["out_time"]) / speed
  124.             self.encode_stats["estimate_remaining"] = remaining_time
  125.         else:
  126.             self.encode_stats["estimate_remaining"] = datetime.timedelta(minutes=0)
  127.  
  128.         if encode_done or self.encode_error:
  129.             status = self.proc.wait()
  130.             end = datetime.datetime.now()
  131.             self.elapsed = end - self.start
  132.             self.stderr.write(f"Encoding finished in {str(self.elapsed)}")
  133.             os.close(self.pipe_write)
  134.             os.close(self.pipe_read)
  135.             self.stderr.close()
  136.         return encode_done or self.encode_error
  137.  
  138.     def cancel(self):
  139.         self.proc.terminate()
  140.         self.proc.wait()
  141.         os.remove(self.out_filename)
  142.         os.remove(self.log_filename)
  143.  
  144. class EncodingManager:
  145.     def __init__(self, all_files, parallel_encodes, main_widget, todo_list, active_list, completed_list):
  146.         self.parallel_encodes = parallel_encodes
  147.         self.main_widget = main_widget
  148.         self.todo_list = todo_list
  149.         self.active_list = active_list
  150.         self.completed_list = completed_list
  151.         self.active_encodes = {}
  152.  
  153.         self.videos = deque()
  154.         for filename, out_filename in all_files:
  155.             try:
  156.                 probe = ffmpeg.probe(filename)
  157.             except ffmpeg.Error as e:
  158.                 continue
  159.  
  160.             if len([s for s in probe["streams"] if s["codec_type"] == "video"]) == 0:
  161.                 continue
  162.  
  163.             duration = float(probe["format"]["duration"])
  164.             seconds = int(duration)
  165.             milliseconds = int((duration - seconds) * 1000)
  166.             duration = datetime.timedelta(seconds=seconds, milliseconds=milliseconds)
  167.  
  168.             video_stream = [s for s in probe["streams"] if s["codec_type"] == "video"][0]
  169.             bitrate = get_video_bitrate(probe)
  170.  
  171.             source_file_ui = urwid.Pile([
  172.                 urwid.Text(filename),
  173.                 urwid.Text(f"Resolution: {video_stream['width']}x{video_stream['height']}\n" +
  174.                     f"Length: {duration}\n" +
  175.                     f"Bitrate: {bitrate}kbits/s"),
  176.                 urwid.Divider("-")
  177.             ])
  178.  
  179.             self.todo_list.body.append(source_file_ui)
  180.             self.videos.append((filename, out_filename))
  181.  
  182.     def monitor_encoding(self):
  183.         self.check_task_completion()
  184.  
  185.         # Start more encodes if we're able to
  186.         if len(self.videos) > 0 and len(self.active_encodes) < self.parallel_encodes:
  187.             self.todo_list.body.pop(0)
  188.             filename, out_filename = self.videos.popleft()
  189.             self.active_encodes[filename] = EncodingTask(filename, out_filename)
  190.             active_encode_ui = urwid.Pile([
  191.                 urwid.Text(filename),
  192.                 urwid.Text(""),
  193.                 urwid.ProgressBar("normal", "complete"),
  194.                 urwid.Text(f"Output: {self.active_encodes[filename].out_filename}"),
  195.                 urwid.Divider("-")])
  196.             self.active_list.body.append(active_encode_ui)
  197.  
  198.         total_fps = 0
  199.         for k, enc in self.active_encodes.items():
  200.             if "fps" in enc.encode_stats:
  201.                 total_fps += enc.encode_stats["fps"]
  202.         self.main_widget.footer.set_text(f"Todo: {len(self.videos)}. Total FPS: {round(total_fps, 2)}. ESC to Cancel/Quit")
  203.  
  204.         # Check which UI column is selected and style the title to indicate it
  205.         columns = self.main_widget.body
  206.         for c in columns.contents:
  207.             title = c[0].title_widget
  208.             if c[0] == columns.focus:
  209.                 title.set_text(("selected_column", title.text))
  210.             else:
  211.                 title.set_text(("default_text", title.text))
  212.  
  213.     def check_task_completion(self):
  214.         complete = []
  215.         for k, enc in self.active_encodes.items():
  216.             ui = [x for x in self.active_list.body if x.contents[0][0].text == k][0]
  217.             if enc.is_complete():
  218.                 complete.append(k)
  219.                 ui.contents[1][0].set_text(f"Resolution: {enc.width}x{enc.height}\n" +
  220.                         f"Bitrate: {enc.encode_stats['bitrate']}\n" +
  221.                         f"FPS: {enc.encode_stats['fps']}\n" +
  222.                         f"Speed: {enc.encode_stats['speed']}\n" +
  223.                         f"Elapsed: {str(enc.elapsed)}")
  224.                 if enc.encode_error:
  225.                     ui.contents[1][0].set_text(ui.contents[1][0].text + "\nEncoding Failed! Check log")
  226.                 else:
  227.                     ui.contents[2][0].set_completion(100)
  228.                 self.completed_list.body.append(ui)
  229.                 self.active_list.body = [x for x in self.active_list.body if x.contents[0][0].text != k]
  230.             else:
  231.                 ui.contents[1][0].set_text(f"Resolution: {enc.width}x{enc.height}\n" +
  232.                         f"Bitrate: {enc.encode_stats['bitrate']}\n" +
  233.                         f"FPS: {enc.encode_stats['fps']}\n" +
  234.                         f"Speed: {enc.encode_stats['speed']}\n" +
  235.                         f"Est. Remaining: {str(enc.encode_stats['estimate_remaining'])}")
  236.                 ui.contents[2][0].set_completion(enc.encode_stats["percent_done"])
  237.         for k in complete:
  238.             del self.active_encodes[k]
  239.  
  240.     def cancel_active_encodes(self):
  241.         for k, enc in self.active_encodes.items():
  242.             enc.cancel()
  243.  
  244. manager = None
  245.  
  246. def quit_on_escape(key):
  247.     if key == "esc":
  248.         manager.cancel_active_encodes()
  249.         raise urwid.ExitMainLoop()
  250.  
  251. def monitor_encoding(loop, man):
  252.     man.monitor_encoding()
  253.     loop.set_alarm_in(0.5, monitor_encoding, user_data=man)
  254.  
  255. if __name__ == "__main__":
  256.     args = docopt.docopt(USAGE)
  257.  
  258.     print("Collecting input video list...")
  259.     all_files = []
  260.     for it in args["<items>"]:
  261.         if os.path.isdir(it):
  262.             for path, dirs, files in os.walk(it):
  263.                 if output_dir in path:
  264.                     continue
  265.                 for f in files:
  266.                     filename = os.path.join(path, f)
  267.                     out_filename = os.path.join(output_dir, os.path.splitext(os.path.relpath(filename, it))[0] + ".mp4")
  268.                     all_files.append((filename, out_filename))
  269.         else:
  270.             out_filename = os.path.join(output_dir, os.path.splitext(it)[0] + ".mp4")
  271.             all_files.append((it, out_filename))
  272.  
  273.     parallel_encodes = 1
  274.     if args["<parallel_encodes>"]:
  275.         parallel_encodes = int(args["<parallel_encodes>"])
  276.  
  277.     # Setup our UI
  278.     todo_list = urwid.ListBox(urwid.SimpleFocusListWalker([]))
  279.     active_list = urwid.ListBox(urwid.SimpleFocusListWalker([]))
  280.     completed_list = urwid.ListBox(urwid.SimpleFocusListWalker([]))
  281.     columns = urwid.Columns([urwid.LineBox(todo_list, title="Todo"),
  282.         urwid.LineBox(active_list, title="Active"),
  283.         urwid.LineBox(completed_list, title="Completed")])
  284.     frame = urwid.Frame(columns,
  285.             header=urwid.Text("FFmpeg Batch Encoding Dashboard"),
  286.             footer=urwid.Text(""))
  287.  
  288.     palette = [
  289.         ("normal", "black", "light gray"),
  290.         ("complete", "black", "dark green"),
  291.         ("selected_column", "black", "white"),
  292.         ("default_text", "white", "black")
  293.     ]
  294.  
  295.     manager = EncodingManager(all_files, parallel_encodes, frame, todo_list, active_list, completed_list)
  296.  
  297.     loop = urwid.MainLoop(frame,
  298.             unhandled_input=quit_on_escape,
  299.             palette=palette)
  300.     loop.set_alarm_in(0.5, monitor_encoding, user_data=manager)
  301.     loop.run()
  302.  
  303.  
Advertisement
Comments
Add Comment
Please, Sign In to add comment
Advertisement