从视频中提取光流 UCF-101

美其名曰:站在巨人的肩膀上

稍微修改,输出变成我想要的样子

来源------->https://github.com/wadhwasahil/Video-Classification-2-Stream-CNN

# main.py

import numpy as np
import optical_flow_prep as ofp
import sys,os
import pickle
import scandir
import gc

def writeOF():

    root = "F:\Dataset/UCF-101"
    w=224
    h=224
    c=0
    data={}

    for path, subdsirs, files in scandir.walk(root):
        for filename in files:
            count=ofp.writeOpticalFlow(path,filename,w,h,c)
            if count:
                data[filename]=count
            print(filename)
            c+=1
            with open("done.txt", "a") as myfile:
                myfile.write(filename+'-'+str(c)+'\n')

    with open('D:\MyDrivers\software/anaconda\wokspace\Video-Classification-2-Stream-CNN-master\dataset/frame_count.pickle','wb') as f:
        pickle.dump(data,f)


def data_prep():
    print('Starting with data prep')
    with open('D:\MyDrivers\software/anaconda\wokspace\Video-Classification-2-Stream-CNN-master/dataset/frame_count.pickle','rb') as f1:
        frame_count=pickle.load(f1)
    with open('D:\MyDrivers\software/anaconda\wokspace\Video-Classification-2-Stream-CNN-master/dataset/merged_data.pickle','rb') as f2:
        merged_data=pickle.load(f2)
    print('Loaded dictionary')
    root = 'F:\Dataset\Optical-Flow'
    path = os.path.join(root, '')
    data={}
    misplaced_data=[]
    count=0
    for path, subdirs, files in scandir.walk(root):
        for filename in files:
            print(filename + '  ' + str(count))
            count+=1
            try:
                vidname=filename.split('_',1)[1].split('.')[0]
                fc=frame_count[vidname]


                for i,j in enumerate(merged_data[vidname]):
                    if j:
                        index=i
                        break
                for i in range(1,(fc/50)+1):
                    data[vidname+'@'+str(i)]=index+1
            except:
                misplaced_data.append(filename)

    print('Writing final training dictionary')
    with open('D:\MyDrivers\software/anaconda\wokspace\Video-Classification-2-Stream-CNN-master/dataset/temporal_train_data.pickle','wb') as f3:
        pickle.dump(data,f3)

    print('Writing misplaced videos')
    with open('D:\MyDrivers\software/anaconda\wokspace\Video-Classification-2-Stream-CNN-master/dataset/misplaced_data.pickle','wb') as f4:
        pickle.dump(misplaced_data,f4   )

if __name__ == "__main__":
    writeOF()
    gc.collect()
    data_prep()

# optical_flow_prep.py

import cv2
import numpy as np
import pickle
from PIL import Image
import os
import gc

def stackOpticalFlow(blocks,temporal_train_data,img_rows,img_cols):
   firstTime=1

   try:
      firstTimeOuter=1
      for block in blocks:
         fx = []
         fy = []
         filename,blockNo=block.split('@')
         path = 'F:\Dataset\Optical-Flow/'#+filename.split('.')[0] + '/'
         blockNo=int(blockNo)

         for i in range((blockNo*10)-9,(blockNo*10)+1):
            # print path+'/'+'h'+str(i)+'_'+str(filename)+'.jpg'
            imgH=Image.open(path+'/'+'h'+str(i*5)+'_'+str(filename)+'.jpg')
            imgV=Image.open(path+'/'+'v'+str(i*5)+'_'+str(filename)+'.jpg')
            imgH=imgH.resize((img_rows,img_cols))
            imgV=imgV.resize((img_rows,img_cols))
            fx.append(imgH)
            fy.append(imgV)
         flowX = np.dstack((fx[0],fx[1],fx[2],fx[3],fx[4],fx[5],fx[6],fx[7],fx[8],fx[9]))
         flowY = np.dstack((fy[0],fy[1],fy[2],fy[3],fy[4],fy[5],fy[6],fy[7],fy[8],fy[9]))
         inp = np.dstack((flowX,flowY))
         inp = np.expand_dims(inp, axis=0)
         if not firstTime:
            inputVec = np.concatenate((inputVec,inp))
            labels=np.append(labels,temporal_train_data[block]-1)
         else:
            inputVec = inp
            labels=np.array(temporal_train_data[block]-1)
            firstTime = 0

      inputVec=np.rollaxis(inputVec,3,1)
      inputVec=inputVec.astype('float16',copy=False)
      labels=labels.astype('int',copy=False)
      gc.collect()

      return (inputVec,labels)
   except:
      return (None,None)


def writeOpticalFlow(path,filename,w,h,c):
   count=0
   try:
      cap = cv2.VideoCapture(path+'/'+filename)
      ret, frame1 = cap.read()

      if frame1 is None:
         return count

      frame1 = cv2.resize(frame1, (w,h))
      prvs = cv2.cvtColor(frame1,cv2.COLOR_BGR2GRAY)

      folder = 'F:/Dataset/Optical-Flow/'+(filename.split('.')[0]).split('_')[1] + '/'
      if not os.path.exists(folder):
         dir = os.path.dirname(folder)
         os.mkdir(dir)
      folder1 = folder + filename.split('.')[0] + '/'
      if not os.path.exists(folder1):
         dir = os.path.dirname(folder1)
         os.mkdir(dir)
      while(1):
         ret, frame2 = cap.read()

         if frame2 is None:
            break
         count+=1
         if count%5==0:
            print (filename+':' +str(c)+'-'+str(count))

            frame2 = cv2.resize(frame2, (w,h))
            next = cv2.cvtColor(frame2,cv2.COLOR_BGR2GRAY)

            flow = cv2.calcOpticalFlowFarneback(prvs,next, None, 0.5, 3, 15, 3, 5, 1.2, 0)

            horz = cv2.normalize(flow[...,0], None, 0, 255, cv2.NORM_MINMAX)
            vert = cv2.normalize(flow[...,1], None, 0, 255, cv2.NORM_MINMAX)
            horz = horz.astype('uint8')
            vert = vert.astype('uint8')

            cv2.imwrite(folder1+'frame_'+str(int(count/5))+'_'+'x'+'.jpg',horz,[int(cv2.IMWRITE_JPEG_QUALITY), 90])
            cv2.imwrite(folder1+'frame_'+str(int(count/5))+'_'+'y'+'.jpg',vert,[int(cv2.IMWRITE_JPEG_QUALITY), 90])
            
            prvs = next

      cap.release()
      cv2.destroyAllWindows()
      return count
   except Exception as e:
      print(e)
      return count

猜你喜欢

转载自blog.csdn.net/weixin_40642281/article/details/81201197