Split theta z1 movie file for meshroom

split_theta_mov.py

# -*- coding: utf-8 -*-
"""
引数にあるMP4ファイルから切りだす。
出力フォルダは  [入力ファイル名]_rig
同時に処理されたファイルはファイル名が被らないようにしてあるので、
各rigフォルダのコピーしてまとめてmeshroomに突っ込む事ができる。

@author: akira_you
"""

SKIP=20 #SKIPフレームに1枚静止画として切りだす
is_rig_cam=True #rigカメラとして出力するか

R_Rate=0.8 # マスクサイズ：有効画像領域半径(1=全部使う)
G_Rate=0.1 # マスクの縁を滑らかにする距離：ガウスぼかしの長さ


import glob
import sys
import os
import cv2 
import pyexiv2
import piexif
import numpy as np
import rawpy
import imageio
import re
maxFrame=999999
exif_dict={"0th":{},"Exif":{}}

def add_serial(output_file_name,rig_index,im):
    global exif_dict
    exif_dict["0th"][piexif.ImageIFD.ImageWidth] = (im.shape[1], 1)
    exif_dict["0th"][piexif.ImageIFD.ImageLength] = (im.shape[0], 1)
    exif_dict["0th"][piexif.ImageIFD.XResolution] = (300, 1)#ここは適当でOK
    exif_dict["0th"][piexif.ImageIFD.YResolution] = (300, 1)#ここは適当でOK
    exif_dict["0th"][piexif.ImageIFD.ResolutionUnit ] = 2
    exif_dict['Exif'][piexif.ExifIFD.PixelXDimension] = im.shape[1]
    exif_dict['Exif'][piexif.ExifIFD.PixelYDimension] = im.shape[0]
    exif_dict["Exif"][piexif.ExifIFD.FocalLength]=   (257,100)#THETA Z1の場合はこの値
    exif_dict["0th"][piexif.ImageIFD.CameraSerialNumber]=  bytes('00{}'.format(rig_index), 'utf-8')
    exif_dict['Exif'][piexif.ExifIFD.LensSerialNumber] =  bytes('00{}'.format(rig_index), 'utf-8')
    exif_dict['Exif'][piexif.ExifIFD.BodySerialNumber] =  bytes('00{}'.format(rig_index), 'utf-8')
    exif_dict['0th'][piexif.ImageIFD.Model] = bytes('CUTED', 'utf-8')
    exif_dict['0th'][piexif.ImageIFD.Make] = bytes('RICOH', 'utf-8')
    exif_bytes = piexif.dump(exif_dict)
    piexif.insert(exif_bytes, output_file_name)
    

def _add_serial(file,s):        
    img = pyexiv2.Image( file )
    metadata = img.read_exif()
    metadata["Exif.Image.CameraSerialNumber"] = s
    img.modify_exif( metadata )
    img.close()
def iwrite(file,img,s):
    cv2.imwrite(file,img,[int(cv2.IMWRITE_JPEG_QUALITY), 99])
    add_serial(file,s,img)

count=0
def main(file):
    global count
    d=file+"rig"
    try:
        os.mkdir(d)
    except:
        return
    try:
        if is_rig_cam:
            for i in range(2):
                os.mkdir("{}/{}".format(d,i))
    except:
        pass
    mask=None
    cap_file = cv2.VideoCapture(file)
    
    while True:
        ret,img=cap_file.read()
        if not ret: break
        count += 1
        if count %SKIP != 0: continue

        #img=rawpy.imread(file).postprocess()[:,:,::-1]
        print(img.shape)
        simg=np.array_split(img, 2, axis=1)
        for i,f_img in enumerate(simg):
            if i==0:
                f_img=f_img.transpose(1,0,2)[:,::-1]
            else:
                f_img=f_img.transpose(1,0,2)[::-1]
            #180度以上はmeshroomのカメラモデル(fisheye4)で扱えないのでマスクする
            if mask is None:
                mask=np.zeros(f_img.shape,np.float32)
                h=f_img.shape[0]
                w=f_img.shape[1]
                cv2.circle(mask, center=(h // 2, w // 2), radius=int(h/2*R_Rate) , color=(1,1,1), thickness=-1)
                mask=cv2.GaussianBlur(mask,(int(h/2*G_Rate)//2*2+1,)*2,0)
                valid_sum=np.sum(mask)
                
            #マスクを強力な特徴量として捉えられないように平均色でマスク
            f_img=(f_img.astype(np.float32)* mask  +  np.sum(f_img)/valid_sum *(1-mask)).astype(np.uint16)
            new_file="{:06}".format(count)
            if is_rig_cam:
                iwrite(d+"/"+str(i)+"/"+new_file+".jpg",f_img,str(i))
            else:
                num=count
                if i == 1:
                    num = maxFrame - count
                new_file="{:06}".format(num)

                iwrite(d+"/"+new_file+".jpg",f_img,str(i))

for f in sys.argv:
    main(f)