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")

args = parser.parse_args()

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

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, intro_start: int = 0):
        self.src = src
        self.dest = dest

        self.temp_dir = tempfile.TemporaryDirectory()
        self.output_dir = self.temp_dir.name
        self.final = os.path.join(self.output_dir, "final")
        self.vid_folder = os.path.join(self.output_dir, "vid")
        self.video_title = None

        os.makedirs(self.final, exist_ok=True)
        os.makedirs(self.vid_folder, exist_ok=True)
        os.makedirs(self.dest, exist_ok=True)

        if self.installation():
            print("[INFO] Finished downloading")
            all_entries = os.listdir(self.vid_folder)
            files = [
                entry
                for entry in all_entries
                if os.path.isfile(os.path.join(self.vid_folder, entry))
            ]
            print(f"[INFO] Found a video {files}, in {self.vid_folder}")

            if files:
                self.cap = cv2.VideoCapture(os.path.join(self.vid_folder, files[0]))
            else:
                print("[ERR] No files found in the directory.")
                exit()
        else:
            self.cap = cv2.VideoCapture(self.src)
            # Extract the filename without extension as a default title
            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.frame_count = intro_start
        self.extracted_count = 0
        self.previous_frame = None

        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        (
            print(f"[DEBUG] Total number of frames in the video: {self.total_frames}")
            if verbose is True
            else None
        )
        self.pbar_count = 0
        self.pbar = tqdm(total=self.total_frames)

    def installation(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.vid_folder}/%(title)s.%(ext)s",
                "quiet": True,  # Suppress output
                "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 _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):
        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("[INFO] Analyzing video...")
                )

            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_count += 1
            self.pbar.update()

        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(key=natural_keys)

        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.final, 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.final, output_filename)
            self.combine_imgs(
                last_image_path, blank_image_path, output_path, mode="vertical"
            )

            os.remove(blank_image_path)

        self.pbar_count += 1
        self.pbar.update()

    def run(self):
        self.analyze_frame()
        self.combine_in_pairs()
        self.convert_to_pdf()

    def scan_folder(self, src):
        files_in_folder = []

        for root, dirs, files in os.walk(src):
            for name in files:
                if name.endswith((".png", ".jpg", ".jpeg", ".mp4", ".webm")):
                    files_in_folder.append(os.path.join(root, name))

        files_in_folder.sort(key=natural_keys)
        return files_in_folder

    def convert_to_pdf(self):
        imgs = self.scan_folder(self.final)
        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))

        self.pbar.set_description(f"[INFO] Saved file to {pdf_path}")
        self.pbar_count += 1
        self.pbar.update()

    def __del__(self):
        self.temp_dir.cleanup()


def atoi(text):
    return int(text) if text.isdigit() else text


def natural_keys(text):
    return [atoi(c) for c in re.split(r"(\d+)", text)]


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