File size: 1,661 Bytes
8b09391 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import os
import cv2
from tqdm import tqdm
class ImageCreator:
def __init__(self, filename, imgs_dir, image_start=0, image_end=0, pbar=True):
"""
:param str filename: The name of the video's filename.
:param str imgs_dir: The directory where to store the image files.
:param int image_start: The first image to be extracted.
:param int image_end: The last image to be extracted, 0 if full video.
:param bool pbar: Whether to display a progress bar.
"""
self.filename = filename
self.imgs_dir = imgs_dir
self.image_start = image_start
self.image_end = image_end
self.pbar = pbar
if not os.path.exists(imgs_dir):
os.makedirs(imgs_dir)
def get_images(self):
vid = cv2.VideoCapture(self.filename)
total_frames = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
success, image = vid.read()
count = 0
if self.image_end == 0:
self.image_end = total_frames
zfill_max = len(str(total_frames))
ok_count = 0
print("Writing images...")
if self.pbar:
pb = tqdm(total=total_frames)
while success:
if count >= self.image_start and count <= self.image_end:
cv2.imwrite(
f"{self.imgs_dir}/frame_{str(ok_count).zfill(zfill_max)}.png", image
)
ok_count += 1
success, image = vid.read()
if self.pbar:
pb.update(1)
count += 1
if self.pbar:
pb.close()
print("Wrote {} image files.".format(ok_count))
|