#!/usr/bin/python3 import numpy as np import cv2 import os import openpiv.tools import openpiv.validation import openpiv.process import openpiv.filters import math from datetime import datetime import multiprocessing CORES=48 #Function accepts directory path as a string variable, creates directory if does not exist, passes to removeDirectory if it does def establishDirectory(directoryPath): if os.path.exists(directoryPath): removeDirectory(directoryPath) os.mkdir(directoryPath) return #Function accepts directory path as a string variable, erases directory and files. def removeDirectory(directoryPath): fileList = os.listdir(directoryPath) for fileName in fileList: os.remove(directoryPath + "/" + fileName) os.rmdir(directoryPath) return ##generalize file for use in other scripts def extractFrames(cap,frameOutputPath='ProcessedFrames_Out'): print("Capturing video.") establishDirectory(frameOutputPath) #Capture video and find frame count, frame rate, and time between frames. frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) frameRate = cap.get(cv2.CAP_PROP_FPS) timeBetweenFrames = (1.000/frameRate) print(str(frameCount) + " frames total.") #Progress indicator print("Beginning loop to read frames, convert to grayscale, and save.") #Loop for image processing video frames startTime=datetime.now() #if False: for i in range(0, frameCount): #for i in range(0, 99): #Read frames, convert to grayscale, and save to directory ret, frame = cap.read() gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) cv2.imwrite(frameOutputPath + "/frame"+str(i)+".png", gray) #Progress indicator. print("frame "+str(i)+" of "+str(frameCount)) #Close video after all frames are stored def processFrameRange(frameOutputPath,pivOutputPath,maskRemoveOutputPath,timeBetweenFrames,start,e): frame_a = openpiv.tools.imread(frameOutputPath + "/frame"+str(start)+".png") startTime=datetime.now() for i in range(start, e): #Establish counter to read sequential frames j = i + 1 #Read current frame and next frame into numpy arrays frame_b = openpiv.tools.imread(frameOutputPath + "/frame"+str(j)+".png") u, v, sig2noise = openpiv.process.extended_search_area_piv(frame_a.astype(np.int32), frame_b.astype(np.int32), window_size=32, overlap=16, dt=timeBetweenFrames, search_area_size=48, subpixel_method='gaussian', sig2noise_method='peak2peak') u, v, mask = openpiv.validation.sig2noise_val(u, v, sig2noise, threshold=1.3) u, v = openpiv.filters.replace_outliers(u, v, method='localmean', max_iter=10, kernel_size=2) #Get window centers coordinates and save velocity field information x, y = openpiv.process.get_coordinates(image_size=frame_a.shape, window_size=32, overlap=16) openpiv.tools.save(x, y, u, v, mask, pivOutputPath + "/piv"+str(i)+".txt") #Open vector information files to begin deleting invalid vectors f = open(pivOutputPath + "/piv"+str(i)+".txt") s = "none" l = [] #Determine which entries have a false mask, indicating a valid vector while s != "": s = f.readline() snip = s[-7:-1] if snip == '0.0000': l.append(s) f.close() #Create new file and write valid vectors. f = open(maskRemoveOutputPath + "/maskRemoved"+str(i)+".txt", 'w') for x in l: f.write(x) f.close() frame_a=frame_b if start==0: print(str((i + 1)*CORES) + " of " + str(e*CORES) + " frames processed and saved. (Time Remaining: "+str((datetime.now()-startTime)/(i+1)*(e-i))+")") def processFrames(frameRate,frameCount,frameOutputPath='ProcessedFrames_Out', pivOutputPath='OpenPIV_Out',maskRemoveOutputPath='Mask_Out' ): #Create directories timeBetweenFrames = (1.0/frameRate) establishDirectory(pivOutputPath) establishDirectory(maskRemoveOutputPath) #Progress indicator #Progress indicator print("Beginning PIV processing.") startTime=datetime.now() #Loop through stored frames, OpenPIV processing. threads=[] divCount=int(math.ceil(frameCount/CORES)) print(frameCount) for i in range(0, CORES): t=multiprocessing.Process(target=processFrameRange,args=(frameOutputPath, pivOutputPath,maskRemoveOutputPath,timeBetweenFrames,i*divCount,min((i+1)*divCount,frameCount-1))) t.start() threads.append(t) count=0 for t in threads: t.join() count+=1 i=count print("joined thread "+str(i)) def video_piv(videoInputPath, frameOutputPath='ProcessedFrames_Out', pivOutputPath='OpenPIV_Out',maskRemoveOutputPath='Mask_Out' ): #Establish directories for file input and output. cap = cv2.VideoCapture(videoInputPath) extractFrames(cap,frameOutputPath) frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) frameRate = cap.get(cv2.CAP_PROP_FPS) cap.release() processFrames(frameRate,frameCount,frameOutputPath,pivOutputPath,maskRemoveOutputPath) #Progress indicator #print("PIV processing complete.") return True if __name__ == '__main__': from argparse import * cwd=os.getcwd() parser=ArgumentParser(description="Tool to obtain wave velocity from video footage. You can use this from another python script by calling video_piv(videoInputPath, frameOutputPath=/'ProcessedFrames_Out/', pivOutputPath=/'OpenPIV_Out/',maskRemoveOutputPath=/'Mask_Out/' )") parser.add_argument('FILE',type=str, help="File to process", default="",nargs='?') parser.add_argument('-f',dest="frame",type=str, help="Target dir to export frames to, defaults to "+cwd+"/ProcessedFrames_Out") parser.add_argument('-p',dest="piv",type=str, help="Target dir to export PIV data to, defaults to "+cwd+"/OpenPIV_Out") parser.add_argument('-m',dest="mask",type=str, help="Target dir to export image mask to, defaults to "+cwd+"/Mask_Out") parser.add_argument('-d',help="Debug by using default internal values",action='store_true') args = parser.parse_args() videoInputPath = cwd+r'/DJI_0001.MOV' frameOutputPath = args.frame pivOutputPath = args.piv maskRemoveOutputPath =args.mask if not args.d: if args.FILE=="": videoInputPath=input("Video to process: ") inputV=input("Path to export frames (leave blank for default): ") if not inputV=="": frameOutputPath = inputV inputV=input("Path to PIV data (leave blank for default): ") if not inputV=="": pivOutputPath = inputV inputV=input("Path to Culled Piv Data (leave blank for default): ") if not inputV=="": maskRemoveOutputPath=inputV else: videoInputPath=args.FILE print(videoInputPath) video_piv(videoInputPath,frameOutputPath,pivOutputPath,maskRemoveOutputPath)