一些CV的小工具代码

切Video为images

"""
Convert each frame in a anime video to 256 x 256 images
"""
import cv2
import os
import numpy as np
import argparse
from moviepy.video.io.VideoFileClip import VideoFileClip
from tqdm import tqdm

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--video-path', type=str, default='/home/ubuntu/Downloads/kimetsu_yaiba.mp4')
    parser.add_argument('--save-path', type=str, default='./script/test_crop')
    parser.add_argument('--max-image', type=int, default=1800)
    parser.add_argument('--start', type=int, default=0)
    parser.add_argument('--end', type=int, default=0)
    parser.add_argument('--image-size', type=int, default=256)

    return parser.parse_args()


class VideoConverter:
    def __init__(self, video_path, save_dir, max_image=1600, start=0, end=0, image_size=256):
        os.makedirs(save_dir, exist_ok=True)

        assert os.path.isfile(video_path), f'{video_path} must be a video file'

        self.save_dir = save_dir
        self.video_path = video_path
        self.max_image = max_image
        self.start = start
        self.end = end or None
        self.counter = 0
        self.image_size = image_size

    def crop_and_save(self, image):
        '''
        Crop a large image into smaller images of given size
        @Returns:
            - True if counter reach max_image
        '''

        height, width, _ = image.shape

        # Image can be divided into rows * cols sub-images
        rows = height // self.image_size
        cols = width // self.image_size

        for r in range(rows):
            for c in range(cols):
                start_x = r * self.image_size
                end_x = start_x + self.image_size
                start_y = c * self.image_size
                end_y = start_y + self.image_size
                # print(f'x: [{start_x}:{end_x}], y: [{start_y}:{end_y}]')
                sub_im = image[start_x: end_x, start_y: end_y, :]

                if np.std(sub_im) > 25.0:
                    save_path = os.path.join(self.save_dir, f'{self.counter}.jpg')
                    cv2.imwrite(save_path, sub_im)
                    self.counter += 1

                if self.counter == self.max_image:
                    return True

        return False

    def process(self):
        '''
        Process video
        '''
        video_clip = VideoFileClip(self.video_path)
        if self.start or self.end:
            video_clip = video_clip.subclip(self.start, self.end)

        video_clip = video_clip.set_fps(video_clip.fps // 10)

        total_frames = round(video_clip.fps * video_clip.duration)
        print(f'Processing video {self.video_path}, {total_frames} frames, size: {video_clip.size}')

        for frame in tqdm(video_clip.iter_frames()):
            # It's better if we resizing before crop to keep the image looks more 'scene'
            h, w, _ = frame.shape
            aspect_ratio = w / h
            ratio = h / (self.image_size * 2)
            w /= ratio
            h = w / aspect_ratio

            frame = cv2.resize(frame[...,::-1], (int(h) , int(w)))
            if self.crop_and_save(frame):
                break

        print(f'Saved {self.counter} images to {self.save_dir}')



if __name__ == '__main__':

    args = parse_args()
    converter = VideoConverter(
        args.video_path,
        args.save_path,
        max_image=args.max_image,
        start=args.start,
        end=args.end,
        image_size=args.image_size
        )

    converter.process()