[cv2] video 데이터

def get_frames(path):
    cap = cv2.VideoCapture(path)
    frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    imgs = ()        
    for fidx in range(frames):
        _, img = cap.read()            
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        imgs.append(img)
    
    ret = {f"image{i}":imgs(i) for i in range(1, len(imgs))}
    ret('image') = imgs(0)
    return ret

def aug(transforms, images):
    res = transforms(**images)
    images = torch.zeros((len(images), 3,180, 320))
    images(0, :, :, :) = res("image")
    for i in range(1, len(images)):
        images(i, :, :, :) = res(f"image{i}")
    return images.permute(0,2,3,1), res

def video_save(frame, output_path , file_name , fps = 10):
    frame_array = ()
    size = 0
    for idx , img in enumerate(frame) : 
        height, width, _ = img.shape
        size = (width,height)
        frame_array.append(np.array(img).astype(np.uint8))
    
    out = cv2.VideoWriter(output_path +"/"+ file_name,cv2.VideoWriter_fourcc(*'DIVX'), fps, size)
    for i in range(len(frame_array)):
        # writing to a image array
        out.write(frame_array(i))
    out.release()

==========================실제 run 블럭======================================
from tqdm import tqdm
from albumentations.pytorch.transforms import ToTensorV2
transforms = A.Compose((
    A.Resize(180,320, p = 1),
    A.ISONoise(p = 0.5),
    A.HorizontalFlip(p = 0.5),
    A.ImageCompression(p = 0.7 , quality_lower=80, quality_upper=100),
    A.ColorJitter(brightness = (0.4,0.7) ,p = 0.6),
    A.Affine( 
        scale  = (0.9,1.2),
        rotate = (-10,10),
        p = 0.8),  
    ToTensorV2(transpose_mask=True)
    ), p=1, additional_targets={f"image{i}":"image" for i in range(1, 50)}) # 

frames = get_frames("./train/TRAIN_0001.mp4")
frames,res = aug(transforms, frames)

output_path = "./new_train"

# 저장할 폴더 생성
if os.path.isdir(output_path):
    print("있음")
else:
    print('없음')
    os.mkdir(output_path)

for i in tqdm(list(df.video_path)):
    # 기존 영상 프레임 가져오기
    frames = get_frames(i)
    # 데이터 변화
    frames,res = aug(transforms, frames)

    # 저장
    video_save(frames, output_path ,"new1_"+i.split('/')(2) , fps = 10)