diff --git a/anifetch.py b/anifetch.py index beefdca..16f5ca2 100644 --- a/anifetch.py +++ b/anifetch.py @@ -6,7 +6,9 @@ import subprocess import sys import time - +import threading +from PIL import Image +import numpy as np def print_verbose(*msg): if args.verbose: @@ -29,6 +31,16 @@ def check_codec_of_file(file:str): codec = str(subprocess.check_output(ffprobe_cmd, text=True).strip()) return codec +def get_width_of_video(file:str): + ffprobe_cmd = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width", "-of", "default=nw=1:nk=1", file] + vid_width = int(subprocess.check_output(ffprobe_cmd, text=True)) + return vid_width + +def get_height_of_video(file:str): + ffprobe_cmd = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=height", "-of", "default=nw=1:nk=1", file] + vid_height = int(subprocess.check_output(ffprobe_cmd, text=True)) + return vid_height + def extract_audio_from_file(file:str, extension): audio_file = BASE_PATH / f"output_audio.{extension}" extract_cmd = ["ffmpeg", "-i", file, "-y", "-vn", "-c:a", "copy", "-loglevel","quiet", audio_file] @@ -159,14 +171,10 @@ def check_chroma_flag(): codec = check_codec_of_file(args.filename) ext = get_ext_from_codec(codec) args.sound_saved_path = str(BASE_PATH / f"output_audio.{ext}") - if args.chroma_flag_given: if args.chroma.startswith("#"): sys.exit("Color for hex code starts with an '0x'! Not a '#'") - - - # check cache old_filename = "" should_update = False @@ -207,48 +215,10 @@ def check_chroma_flag(): WIDTH = args.width HEIGHT = args.height - # put cached frames here frames: list[str] = [] -# cache is invalid, re-render -if should_update: - print_verbose("SHOULD RENDER WITH CHAFA") - - # delete all old frames - shutil.rmtree(BASE_PATH / "video") - os.mkdir(BASE_PATH / "video") - - stdout = None if args.verbose else subprocess.DEVNULL - stderr = None if args.verbose else subprocess.STDOUT - - if args.chroma_flag_given: - subprocess.call( - [ - "ffmpeg", - "-i", - f"{args.filename}", - "-vf", - f"fps={args.framerate},format=rgba,chromakey={args.chroma}", - str(BASE_PATH / "video/%05d.png"), - ], - stdout=stdout, - stderr=stderr, - ) - else: - subprocess.call( - [ - "ffmpeg", - "-i", - f"{args.filename}", - "-vf", - f"fps={args.framerate},format=rgba", - str(BASE_PATH / "video/%05d.png"), - ], - stdout=stdout, - stderr=stderr, - ) - +def get_sound(): print_verbose(args.sound_flag_given) if args.sound_flag_given: @@ -269,6 +239,104 @@ def check_chroma_flag(): print_verbose(args.sound_saved_path) +thread_sound = threading.Thread(target=get_sound) + +def chafa_process(f, ffmpeg_frame): + WIDTH = args.width + HEIGHT = args.height + + chafa_args = args.chafa_arguments.strip() + chafa_args += " --format symbols" # Fixes https://github.com/Notenlish/anifetch/issues/1 + + path = BASE_PATH / "video" / f + im = Image.fromarray(ffmpeg_frame) + im.save(path, compress_level=0) + chafa_cmd = [ + "chafa", + *chafa_args.split(" "), + # "--color-space=rgb", + f"--size={WIDTH}x{HEIGHT}", + path.as_posix(), + ] + frame = subprocess.check_output( + chafa_cmd, + text=True, + ) + + with open((BASE_PATH / "output" / f).with_suffix(".txt"), "w") as file: + file.write(frame) + + # if wanted aspect ratio doesnt match source, chafa makes width as high as it can, and adjusts height accordingly. + # AKA: even if I specify 40x20, chafa might give me 40x11 or something like that. + HEIGHT = len(frame.splitlines()) + frames.append(frame) # dont question this, I need frames to have at least a single item + +def ffmpeg_process(): + threads = [] + stderr = None if args.verbose else subprocess.PIPE + vid_width = get_width_of_video(args.filename) + vid_height = get_height_of_video(args.filename) + if args.chroma_flag_given: + ffmpeg_cmd = [ + "ffmpeg", + "-i", + f"{args.filename}", + "-f", + "image2pipe", + "-vf", + f"fps={args.framerate},chromakey={args.chroma}", + "-pix_fmt", + "rgb24", + "-vcodec", + "rawvideo", + "-", + ] + else: + ffmpeg_cmd = [ + "ffmpeg", + "-i", + f"{args.filename}", + "-f", + "image2pipe", + "-vf", + f"fps={args.framerate}", + "-pix_fmt", + "rgb24", + "-vcodec", + "rawvideo", + "-", + ] + print_verbose(vid_width) + print_verbose(vid_height) + proc = subprocess.Popen(ffmpeg_cmd, stdout = subprocess.PIPE, stderr = stderr, bufsize=10**8) + frame_size = vid_width * vid_height * 3 + i=1 + while True: + raw_frame = proc.stdout.read(frame_size) + if not raw_frame: + break + ffmpeg_frame = np.frombuffer(raw_frame, np.uint8).reshape((vid_height, vid_width, 3)) + proc.stdout.flush() + f = str(i) + ".png" + thread_chafa = threading.Thread(target=chafa_process, args=(f, ffmpeg_frame )) + thread_chafa.start() + threads.append(thread_chafa) + print_verbose("Launching chafa thread") + i=i+1 + + + for i in enumerate(threads): + thread_chafa.join() + +thread_ffmpeg = threading.Thread(target=ffmpeg_process) + +# cache is invalid, re-render +if should_update: + print_verbose("SHOULD RENDER WITH CHAFA") + + # delete all old frames + shutil.rmtree(BASE_PATH / "video") + os.mkdir(BASE_PATH / "video") # If the new anim frames is shorter than the old one, then in /output there will be both new and old frames. Empty the directory to fix this. shutil.rmtree(BASE_PATH / "output") @@ -276,38 +344,15 @@ def check_chroma_flag(): print_verbose("Emptied the output folder.") - # get the frames - animation_files = os.listdir(BASE_PATH / "video") - animation_files.sort() - for i, f in enumerate(animation_files): - # TODO: REMOVE THIS - #print_verbose(f"- Frame: {f}") - - # f = 00001.png - chafa_args = args.chafa_arguments.strip() - chafa_args += " --format symbols" # Fixes https://github.com/Notenlish/anifetch/issues/1 - - path = BASE_PATH / "video" / f - chafa_cmd = [ - "chafa", - *chafa_args.split(" "), - # "--color-space=rgb", - f"--size={WIDTH}x{HEIGHT}", - path.as_posix(), - ] - frame = subprocess.check_output( - chafa_cmd, - text=True, - ) + thread_ffmpeg.start() - with open((BASE_PATH / "output" / f).with_suffix(".txt"), "w") as file: - file.write(frame) + thread_sound.start() - # if wanted aspect ratio doesnt match source, chafa makes width as high as it can, and adjusts height accordingly. - # AKA: even if I specify 40x20, chafa might give me 40x11 or something like that. - if i == 0: - HEIGHT = len(frame.splitlines()) - frames.append(frame) # dont question this, I need frames to have at least a single item + thread_ffmpeg.join() + + thread_sound.join() + + # get the frames else: # just use cached for filename in os.listdir(BASE_PATH / "output"): diff --git a/flake.nix b/flake.nix index 57353be..abb6e1c 100644 --- a/flake.nix +++ b/flake.nix @@ -24,5 +24,24 @@ overlays = import ./nix/overlays; formatter = forAllSystems (system: nixpkgs.legacyPackages.${system}.alejandra); + devShells = forAllSystems (system: let + pkgs = import nixpkgs {inherit system;}; + in { + default = pkgs.mkShell { + packages = [ + (pkgs.python3.withPackages (python-pkgs: [ + python-pkgs.pillow + python-pkgs.numpy + ])) + pkgs.neofetch + pkgs.fastfetch + pkgs.chafa + pkgs.bc + (pkgs.callPackage ./nix/packages/anifetch.nix {}) + ]; + shellHook = '' + ''; + }; + }); }; } diff --git a/nix/packages/anifetch.nix b/nix/packages/anifetch.nix index 2a085ca..466b5d4 100644 --- a/nix/packages/anifetch.nix +++ b/nix/packages/anifetch.nix @@ -20,6 +20,8 @@ in pkgs.bc pkgs.chafa pkgs.ffmpeg + pkgs.python3Packages.pillow + pkgs.python3Packages.numpy loop ]; preBuild = ''