get_chara_shots.py

import cv2
import argparse
from PIL import Image
import os
import numpy as np
import SSBUBoundingBoxUtil

parser = argparse.ArgumentParser()
parser.add_argument('in_file', type=str)
parser.add_argument('out_dir', type=str)
parser.add_argument('fighter_num', type=int)
args = parser.parse_args()


os.makedirs(args.out_dir, exist_ok=True)

frame_files = None
cap = None
if os.path.isdir(args.in_file):
    frame_files = [ os.path.join(args.in_file, file) for file in os.listdir(args.in_file) ]
else:
    cap = cv2.VideoCapture(args.in_file)

frame_idx = 0
def next_frame():
    global frame_idx

    if frame_files is not None:
        if len(frame_files) <= frame_idx:
            return None
        frame = cv2.imread(frame_files[frame_idx])
    else:
        ret, frame = cap.read()

    frame_idx += 1
    return frame

in_file = os.path.basename(args.in_file)

while True:
    img = next_frame()
    if img is None:
        break

    bboxes = SSBUBoundingBoxUtil.fighters_chara_bbox(args.fighter_num)
    for fighter_idx in range(len(bboxes)):
        # if fighter_idx != 3:
        #     continue
        bbox = bboxes[fighter_idx]
        chara = img[bbox[1]:bbox[1]+bbox[3], bbox[0]:bbox[0]+bbox[2]]

        file = '%s-f%d-%d.png' % (in_file, fighter_idx, frame_idx, )
        path = os.path.join(args.out_dir, file)
        cv2.imwrite(path, chara)

        print(frame_idx, fighter_idx)