[cv2] video 데이터 증량하기

현재 video data를 가지고 분류하는 테스크 대회를 진행하고 있다.

 

이미지 증량은 경험이 있지만, video의 증량은 경험이 전무하기에 조금 애 먹었다.

 

하지만, 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)

 

위 코드는 이미있는 mp4 영상을 albumentation으로 변환하여 다시금 저장하는 코드이다.

 

1. cv2로 기존 mp4영상을 가져와 이미지를 모은다.

    -위 코드에서 딕셔너리로 ret를 만드는 이유는 밑에서 albumentation을 동일하게 적용하기 위한 기준이 필요하여 위와같이 작성하였다.

2. 가져온 프레임 list를 aug함수에 넣어 albumentation에서 지정해준 변환을 수행한다.

3.여기서, 모델에 넣고 돌릴예정이라, ToTensorV2()를 사용해 tensor로 바꿔줄때, HWC > CHW 로 바뀌는걸 인지하고 있어야 한다.

4. cv2 videowriter를 선언해주고, 저장해주면 끝.

 

간단하게 할 수 있다.