import os
import cv2
from tqdm import tqdm

base_dir = '/home/simplew/dataset/sod/DHF1K'

video_dir = os.path.join(base_dir, 'video')

movies = [mov for mov in os.listdir(video_dir) if mov.endswith('.AVI')]
for movie in tqdm(movies):
    image_dir = os.path.join(base_dir , 'annotation', '0' + movie[:-4], 'images')
    os.makedirs(image_dir, exist_ok=True)
    
    # use opencv
    # cap = cv2.VideoCapture(f"{video_dir}/{movie}")
    # numFrames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # for k in range(numFrames):
    #     ret, frame = cap.read()
    #     cv2.imwrite(f"{image_dir}/{k+1:04}.png", frame)
    # cap.release()

    # ubuntu
    command = f"ffmpeg -i {video_dir}/{movie}  {image_dir}/%04d.png"
    os.system(command)