1、预测函数
导入rgb和flow的CBAM模型,进行预测,将二者预测结果平均法求最终预测结果返回。
def predict_action(vid,frame_count):
rgb_model_path = 'checkpoint/rgb_resnet152_cbam_model_best.pth.tar'
flow_model_path = 'checkpoint/flow_resnet152_cbam_model_best.pth.tar'
start_frame = 0
num_categories = 101
rgb_params = torch.load(rgb_model_path ,map_location='cpu')
spatial_net = rgb_resnet_cbam.rgb_resnet152_cbam(pretrained=False, num_classes=101)
spatial_net.load_state_dict({k.replace('module.',''):v for k,v in rgb_params['state_dict'].items()})
spatial_net.eval()
spatial_prediction = SpatialPrediction.VideoSpatialPrediction(vid,spatial_net,num_categories,start_frame,frame_count)
avg_spatial_pred_fc8 = np.mean(spatial_prediction, axis=1)
#找到每列中最大数即为预测结果
flow_params = torch.load(flow_model_path ,map_location='cpu')
temporal_net = flow_resnet_cbam.flow_resnet152_cbam(pretrained=False, num_classes=101)
temporal_net.load_state_dict({k.replace('module.',''):v for k,v in flow_params['state_dict'].items()})
temporal_net.eval()
temporal_prediction = TemporalPrediction.VideoTemporalPrediction(vid,temporal_net,num_categories,start_frame,frame_count)
avg_temporal_pred_fc8 = np.mean(temporal_prediction, axis=1)
pred_index = np.argmax((avg_temporal_pred_fc8 + avg_spatial_pred_fc8)/2)
print('prediction done!')
return pred_index+1
2、index转action类别
按行读取classInd.txt文件,元组形式返回index:action
def find_classes(class_file):
f = open(class_file, "r")
class_ind = dict()
val_list = f.readlines()
for line in val_list:
line_info = line.split(" ")
cl = line_info[0]
action = line_info[1]
class_ind[cl] = action
return class_ind
3、flask搭建
flk = Flask(name)构造一个flask实体,接收静态文件static下的index.html和templates下的js文件。设置上传视频保存路径,提取帧保存路径以及上传视频的格式限制。
URL与函数绑定,’/index’返回一个网页,将本地的index.html发送给浏览器显示网页
'/predict’实现帧提取,动作识别,并将识别结果返回给浏览器,使浏览器实现动态数据变换。
def config():
#获取当前模块名,确定该模块的根目录,接收静态静态文件static和模板文件templates的目录
flk = Flask(__name__)
#session秘钥配置
flk.config['SECRET_KEY'] = 'hard to guess'
return flk
app = config()
print(app)
app.config['JSON_AS_ASCII'] = False
UPLOAD_VID_PATH = 'video/'
GENERATE_PIC_PATH = 'static/generated'
#设置允许上传图片的格式
ALLOWED_EXTENSIONS = set(['avi', 'mp4'])
def allowed_file(filename):
return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS
#表示传递一个网站,“/index”是网站的主目录
#从服务器上获取页面返回给浏览器
@app.route('/index', methods=["GET"])
#index()是传输给route的函数,里面返回值index.html就是显示到网页上的内容
def index():
#创建视图函数,将模板中的index.html渲染后返回
return render_template('index.html')
#从服务器上获取信息后GET,更新URL上的内容POST
@app.route('/predict', methods=['GET', 'POST'])
def generate_pic():
# 从前端获取客户端上传的图片并保存下来
##"pic"对应前端表单name属性
f = request.files['video']
if not (f and allowed_file(f.filename)):
return "请检查上传的视频类型,仅限于avi,mp4"
#将视频保存在UPLOAD_VID_PATH目录下
update_path = os.path.join(UPLOAD_VID_PATH, f.filename)
f.save(update_path)
print('{}:{}{}{}'.format('Video',f.filename,' has saved to ',update_path))
i = random.randint(1,6)
doc = f.filename.split('.')[0]
#获取视频双流
frame_count = extract.extract_rgb(update_path,GENERATE_PIC_PATH,'img_','flow_x_','flow_y_')
#获取视频的rgb帧和flow帧路径
path = GENERATE_PIC_PATH+'/'+doc
#two-steam预测结果
pred_ind = predict_action(path,frame_count)
class_to_ind = find_classes('datasets/ucf101_splits/classInd.txt')
i = random.randint(1,frame_count)
rgb = '{}/{}{:05d}{}'.format(path,'img_',i,'.jpg')
flow_x = '{}/{}{:05d}{}'.format(path,'flow_x',i,'.jpg')
flow_y = '{}/{}{:05d}{}'.format(path,'flow_y',i,'.jpg')
result=dict()
result['path'] = rgb
result['prediction'] = class_to_ind[str(pred_ind)]
if result is not None:
return json.dumps(result)
return json.dumps("NONE")
if __name__ == '__main__':
#开启服务器功能,仅本机可以访问
http_server = WSGIServer(('127.0.0.1', 8177), app)
http_server.serve_forever()
print('server served!')