diff --git a/vid2sheet.py b/vid2sheet.py index fbfb604..74c22f1 100644 --- a/vid2sheet.py +++ b/vid2sheet.py @@ -25,12 +25,20 @@ parser.add_argument("source", type=str, help="source of the file or a YouTube li 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}") @@ -38,58 +46,41 @@ print("[INFO] Verbose enabled") if verbose is True else None class Vid2Sheet: - def __init__(self, src, dest, intro_start: int = 0): + 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.final = os.path.join(self.output_dir, "final") - self.vid_folder = os.path.join(self.output_dir, "vid") + 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.final, exist_ok=True) - os.makedirs(self.vid_folder, exist_ok=True) + os.makedirs(self.combined_img_dir, exist_ok=True) + os.makedirs(self.video_dir, 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.frame_count = frame_starts_at 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) + # self.pbar = None + self.total_frames = 0 - def installation(self): + 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/)" @@ -99,8 +90,8 @@ class Vid2Sheet: 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 + "outtmpl": f"{self.video_dir}/%(title)s.%(ext)s", + "quiet": True, "progress_hooks": [self._hook], } print("[INFO] Attempting to start download...") @@ -112,11 +103,52 @@ class Vid2Sheet: 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: @@ -133,7 +165,11 @@ class Vid2Sheet: ( 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...") + ) + if verbose is True + else self.pbar.set_description( + f"[INFO] Analyzing {self.video_title}" + ) ) else: @@ -145,16 +181,22 @@ class Vid2Sheet: 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.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.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() @@ -197,7 +239,7 @@ class Vid2Sheet: for f in non_hidden_files if os.path.isfile(os.path.join(self.output_dir, f)) ] - images.sort(key=natural_keys) + images.sort() if len(images) % 2 != 0: last_image = images.pop() @@ -208,7 +250,7 @@ class Vid2Sheet: 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) + 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: @@ -221,34 +263,23 @@ class Vid2Sheet: 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) + 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) - 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) + 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 @@ -259,22 +290,12 @@ class Vid2Sheet: 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() + print(f"[INFO] Saved file to {pdf_path}") def __del__(self): + self.pbar.close() # Close the progress bar when done 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()