kenji/vid2sheet
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

🧑💻 made it more readable

Browse Source
This commit is contained in:
sakamoto
2024-08-03 09:52:33 -05:00
parent 0b3b804260
commit a3f1d318b9
1 changed files with 103 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vid2sheet.py
+103 -82
View File
@@ -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()