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