美其名曰:站在巨人的肩膀上
稍微修改,输出变成我想要的样子
来源------->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