vid2sheet/vid2sheet.py

import cv2
import numpy as np
import os
import re
import img2pdf
import tempfile
import yt_dlp
import argparse

from PIL import Image
from tqdm import tqdm

# test
# https://www.youtube.com/watch?v=tyloC0e-Tqk
# ./

# this is so overengineered that the bloat takes up 1/4 of the line

parser = argparse.ArgumentParser(
    description="Converts Video static images based on significant frame changes to sheet music in a form of .pdf file.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)

parser.add_argument("source", type=str, help="source of the file or a YouTube link")
parser.add_argument("destination", type=str, help="destination of the output file")

parser.add_argument("-v", "--verbose", action="store_true", help="enable debug mode")
parser.add_argument(
    "-t",
    "--change-threshold",
    type=int,
    default=12500000,
    help="take a screenshot based on threshold",
)

args = parser.parse_args()

src = args.source
dest = args.destination
verbose = args.verbose
change_threshold = args.change_threshold

print(f"[INFO] The source file is: {src}")
print(f"[INFO] The destination file is: {dest}")
print("[INFO] Verbose enabled") if verbose is True else None


class Vid2Sheet:
    def __init__(self, src, dest, frame_starts_at: int = 0):
        self.src = src
        self.dest = dest

        self.temp_dir = tempfile.TemporaryDirectory()
        self.output_dir = self.temp_dir.name
        self.combined_img_dir = os.path.join(self.output_dir, "combined_img")
        self.video_dir = os.path.join(self.output_dir, "video")
        self.video_title = None

        os.makedirs(self.combined_img_dir, exist_ok=True)
        os.makedirs(self.video_dir, exist_ok=True)
        os.makedirs(self.dest, exist_ok=True)

        self.frame_count = frame_starts_at
        self.extracted_count = 0
        self.previous_frame = None

        # self.pbar = None
        self.total_frames = 0

        if verbose:
            print(f"[DEBUG] temp_dir: {self.temp_dir}")
            print(f"[DEBUG] output_dir: {self.output_dir}")
            print(f"[DEBUG] combined_img_dir: {self.combined_img_dir}")
            print(f"[DEBUG] video_dir: {self.video_dir}")

    def run(self):
        self.check_video()
        self.analyze_frame(change_threshold)
        self.combine_in_pairs()
        self.pbar.close()
        self.convert_to_pdf()

    def install_yt(self):
        youtube_pattern = re.compile(
            r"(https?://)?(www\.)?"
            r"(youtube\.com/watch\?v=|youtu\.be/)"
            r"[a-zA-Z0-9_-]{11}",
            re.IGNORECASE,
        )
        if re.match(youtube_pattern, self.src):
            print("[INFO] Detected YouTube link")
            ydl_opts = {
                "outtmpl": f"{self.video_dir}/%(title)s.%(ext)s",
                "quiet": True,
                "progress_hooks": [self._hook],
            }
            print("[INFO] Attempting to start download...")
            try:
                with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                    ydl.download([self.src])
                return True
            except Exception as e:
                print(f"[ERR] An error occurred: {e}")
                return False

    def check_video(self):
        if self.install_yt():
            print("[INFO] Finished downloading")
            all_entries = os.listdir(self.video_dir)
            files = [
                entry
                for entry in all_entries
                if os.path.isfile(os.path.join(self.video_dir, entry))
            ]
            print(f"[INFO] Found {files}, in {self.video_dir}")

            if files:
                self.cap = cv2.VideoCapture(os.path.join(self.video_dir, files[0]))
            else:
                print(f"[ERR] No files found in the directory {self.video_dir}.")
                exit()
        else:
            self.cap = cv2.VideoCapture(self.src)
            self.video_title = os.path.splitext(os.path.basename(self.src))[0]

        if not self.cap.isOpened():
            print("[ERR] Could not open video.")
            exit()

        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.pbar = tqdm(
            total=self.total_frames,
            desc="Analyzing Frames",
            bar_format="{l_bar}{bar} | {n_fmt}/{total_fmt} frames | {rate_fmt} | {elapsed} elapsed",
        )
        (
            print(f"[DEBUG] Total number of frames in the video: {self.total_frames}")
            if verbose is True
            else None
        )

    def _hook(self, d):
        if d["status"] == "finished":
            self.video_title = d.get("info_dict", {}).get("title", "unknown_title")

    def analyze_frame(self, change_threshold=12500000):
        (
            print(f"[DEBUG] Change threshold is set to {change_threshold}")
            if verbose is True
            else None
        )
        while True:
            ret, current_frame = self.cap.read()
            if not ret:
                break

            self.gray_current = cv2.cvtColor(current_frame, cv2.COLOR_BGR2GRAY)

            if self.previous_frame is None:
                image_path = os.path.join(
                    self.output_dir, f"image_{self.extracted_count:03}.jpg"
                )
                cv2.imwrite(image_path, current_frame)
                self.extracted_count += 1
                (
                    self.pbar.set_description(
                        f"[DEBUG] Start at Frame {self.frame_count}, saved as {image_path}"
                    )
                    if verbose is True
                    else self.pbar.set_description(
                        f"[INFO] Analyzing {self.video_title}"
                    )
                )

            else:
                frame_diff = cv2.absdiff(self.previous_frame, self.gray_current)
                diff_sum = np.sum(frame_diff)

                if diff_sum > change_threshold:
                    image_path = os.path.join(
                        self.output_dir, f"image_{self.extracted_count:03}.jpg"
                    )
                    cv2.imwrite(image_path, current_frame)
                    (
                        self.pbar.set_description(
                            f"[DEBUG] Frame {self.frame_count} changed significantly, saved as {image_path}"
                        )
                        if verbose is True
                        else None
                    )
                    self.extracted_count += 1

            self.previous_frame = self.gray_current
            self.frame_count += 1

            self.pbar.update(1)  # Update the progress bar by 1 for each frame

        if self.pbar.n < self.total_frames:
            self.pbar.update(self.total_frames - self.pbar.n)

        self.cap.release()

    def create_blank_image(self, width, height):
        return Image.new("RGB", (width, height), "white")

    def combine_imgs(self, image_1, image_2, dest, mode="vertical"):
        if isinstance(image_1, str):
            image_1 = Image.open(image_1)
        if isinstance(image_2, str):
            image_2 = Image.open(image_2)

        width_1, height_1 = image_1.size
        width_2, height_2 = image_2.size

        if mode == "horizontal":
            total_width = width_1 + width_2
            max_height = max(height_1, height_2)
            combined_image = Image.new("RGB", (total_width, max_height))
            combined_image.paste(image_1, (0, 0))
            combined_image.paste(image_2, (width_1, 0))

        elif mode == "vertical":
            max_width = max(width_1, width_2)
            total_height = height_1 + height_2
            combined_image = Image.new("RGB", (max_width, total_height))
            combined_image.paste(image_1, (0, 0))
            combined_image.paste(image_2, (0, height_1))

        else:
            raise ValueError("[ERR] Mode must be either 'vertical' or 'horizontal'.")

        combined_image.save(dest)

    def combine_in_pairs(self):
        all_files = os.listdir(self.output_dir)
        non_hidden_files = [f for f in all_files if not f.startswith(".")]
        images = [
            f
            for f in non_hidden_files
            if os.path.isfile(os.path.join(self.output_dir, f))
        ]
        images.sort()

        if len(images) % 2 != 0:
            last_image = images.pop()
        else:
            last_image = None

        for img in range(0, len(images), 2):
            image_1 = os.path.join(self.output_dir, images[img])
            image_2 = os.path.join(self.output_dir, images[img + 1])
            output_filename = f"combined_{img//2 + 1:03}.jpg"
            output_path = os.path.join(self.combined_img_dir, output_filename)
            self.combine_imgs(image_1, image_2, output_path, mode="vertical")

        if last_image:
            last_image_path = os.path.join(self.output_dir, last_image)
            last_image_img = Image.open(last_image_path)
            width, height = last_image_img.size
            blank_image = self.create_blank_image(width, height)

            blank_image_path = os.path.join(self.output_dir, "blank_image.jpg")
            blank_image.save(blank_image_path)

            output_filename = f"combined_{len(images)//2 + 1:03}.jpg"
            output_path = os.path.join(self.combined_img_dir, output_filename)
            self.combine_imgs(
                last_image_path, blank_image_path, output_path, mode="vertical"
            )

            os.remove(blank_image_path)

    def convert_to_pdf(self):
        all_entries = os.listdir(self.combined_img_dir)
        imgs = [
            os.path.join(self.combined_img_dir, entry)
            for entry in all_entries
            if os.path.isfile(os.path.join(self.combined_img_dir, entry))
            and entry.endswith(".jpg")
        ]

        print(f"[DEBUG] converting these imgs: {imgs} to .pdf") if verbose is True else print("[INFO] Converting to pdf...")
        if not imgs:
            print("[ERR] No images found for PDF conversion.")
            return

        file_name = self.video_title if self.video_title else "output"
        pdf_path = os.path.join(self.dest, f"{file_name}.pdf")

        with open(pdf_path, "wb") as f:
            f.write(img2pdf.convert(imgs))

        print(f"[INFO] Saved file to {pdf_path}")

    def __del__(self):
        self.pbar.close()  # Close the progress bar when done
        self.temp_dir.cleanup()

if __name__ == "__main__":
    program = Vid2Sheet(src, dest)
    program.run()