7月31日更新PP_TSMv2模型推理输出

tool/PP_TSMv2_infer.py

@@ -1,29 +1,22 @@
-from export_model import trim_config,get_input_spec
-from predict import parse_file_paths
 import os 
 import os.path as osp
-import paddle
-from paddlevideo.utils import get_config
-from paddlevideo.modeling.builder import build_model
-from paddle.jit import to_static
+from paddlevideo.utils.config import get_config
 from paddle.inference import Config, create_predictor
 from utils import build_inference_helper
-import time
 import warnings
 warnings.filterwarnings("ignore")
 
 
-class PP_TSMv2(object):
+class PP_TSMv2_predict(object):
 
     """PP-TSMv2模型中常用的参数初始化"""
 
-    def __init__(self,use_gpu=True,batch_size=1,ir_optim=True,\
-                 disable_glog=False,save_name=None,enable_mklddn=False,\
+    def __init__(self,use_gpu=True,ir_optim=True,
+                 disable_glog=False,save_name=None,enable_mklddn=False,
                  precision="fp32",gpu_mem=8000,cpu_threads=None):
 
         self.use_gpu = use_gpu                                 #是否使用GPU
         self.cpu_threads = cpu_threads                         #cpu线程数
-        self.batch_size = batch_size
         self.ir_optim = ir_optim                               #是否开启IR优化
         self.disable_glog = disable_glog
         self.gpu_mem = gpu_mem                                 #GPU存储大小
@@ -31,6 +24,25 @@ class PP_TSMv2(object):
         self.precision = precision                             #mfldnn精度
         self.save_name = save_name                             #转化推理模型存放名称
 
+    def parse_file_paths(self,input_path: str) -> list:
+        
+        """
+            获取模型输入数据
+            input_path:模型的输入文件
+        """
+        if osp.isfile(input_path):
+            files = [
+                input_path,
+            ]
+        else:
+            files = os.listdir(input_path)
+            files = [
+                file for file in files
+                if (file.endswith(".avi") or file.endswith(".mp4"))
+            ]
+            files = [osp.join(input_path, file) for file in files]
+        return files
+
 
     def create_paddle_predictor(self,model_f,pretr_p,cfg):
         """
@@ -66,47 +78,6 @@ class PP_TSMv2(object):
         return config,predictor
 
 
-    def exportmodel(self,config,pretr_p,output_p):
-
-        """
-            加载训练后的模型参数，生成可推理预测的模型
-            pretr_p:训练后的参数存放文件
-            output_p:转化为可推理模型的存放路径
-        """
-
-        cfg, model_name = trim_config(get_config(config, overrides=None, show=False))
-
-        print(f"Building model({model_name})...")
-
-        #创建推理模型
-        model = build_model(cfg.MODEL)
-        assert osp.isfile(
-            pretr_p
-        ), f"pretrained params ({pretr_p} is not a file path.)"
-
-        if not os.path.isdir(output_p):
-            os.makedirs(output_p)
-
-        print(f"Loading params from ({pretr_p})...")
-
-        # 加载推理模型参数
-        params = paddle.load(pretr_p)
-        model.set_dict(params)
-
-        model.eval()
-
-        for layer in model.sublayers():
-            if hasattr(layer, "rep") and not getattr(layer, "is_repped"):
-                layer.rep()
-
-        input_spec = get_input_spec(cfg.INFERENCE, model_name)
-
-        #将模型转化为静态图以及模型的保存
-        model = to_static(model, input_spec=input_spec)
-        paddle.jit.save(model,osp.join(output_p, model_name if self.save_name is None else self.save_name))
-        print(f"model ({model_name}) has been already saved in ({output_p}).")
-        return model
-    
 
     def predict(self,config,input_f,batch_size,model_f,params_f):
 
@@ -118,7 +89,7 @@ class PP_TSMv2(object):
             model_f:可推理模型存放的路径+配置文件
             params_f:可推理模型的参数
         """
-
+        result = {}
         cfg = get_config(config, overrides=None, show=False)
         model_name = cfg.model_name
         print(f"Inference model({model_name})...")
@@ -139,7 +110,7 @@ class PP_TSMv2(object):
         for item in output_names:
             output_tensor_list.append(predictor.get_output_handle(item))
 
-        files = parse_file_paths(input_f)
+        files = self.parse_file_paths(input_f)#input_path=input_f
 
         batch_num = batch_size
         for st_idx in range(0, len(files), batch_num):
@@ -158,22 +129,26 @@ class PP_TSMv2(object):
                 batched_outputs.append(output_tensor_list[j].copy_to_cpu())
 
             #输出推理结果
-            InferenceHelper.postprocess(batched_outputs,True)
+            res = InferenceHelper.postprocess(batched_outputs,False,True)     
+        result["video_id"] = res[0]["video_id"]
+        result["topk_class"] = res[0]["topk_class"].tolist()[0]
+        result["topk_scores"] = res[0]["topk_scores"].tolist()[0]
+        
+        return result
+            
 
 
 
 def main():
 
-    config='/home/xznsh/data/PaddleVideo/configs/recognition/pptsm/v2/pptsm_lcnet_k400_16frames_uniform.yaml'           #配置文件地址
-    input_file='/home/xznsh/data/PaddleVideo/data/dataset/video_seg_re_hand'                                            #待推理数据集存放的地址
-    pretrain_params='/home/xznsh/data/PaddleVideo/output/ppTSMv2/ppTSMv2_best.pdparams'                                 #训练后模型参数文件存放
-    output_path='/home/xznsh/data/PaddleVideo/inference/infer1'                                                         #推理模型存放地址
-    model_file='/home/xznsh/data/PaddleVideo/inference/infer1/ppTSMv2.pdmodel'                                          #推理模型存放地址
-    params_file='/home/xznsh/data/PaddleVideo/inference/infer1/ppTSMv2.pdiparams'                                       #推理模型参数存放地址
-    batch_size= 1
-    PP_TSMv2().exportmodel(config,pretrain_params,output_path)                                                          #输出推理模型
-    time.sleep(2)
-    PP_TSMv2().predict(config,input_file,batch_size,model_file,params_file)                                             #推理模型推理、预测
+    config='D:/download/PaddleVideo1/configs/recognition/pptsm/v2/pptsm_lcnet_k400_16frames_uniform.yaml'           #配置文件地址
+    input_file='D:/download/PaddleVideo1/data/dataset/video_seg_no_hand/test02_84.avi'                              #待推理数据集存放的地址
+    model_file='D:/download/PaddleVideo1/output/ppTSMv2.pdmodel'                                          #推理模型存放地址
+    params_file='D:/download/PaddleVideo1/output/ppTSMv2.pdiparams'                                       #推理模型参数存放地址
+    batch_size= 1                                                          #输出推理模型
+    t = PP_TSMv2_predict().predict(config,input_file,batch_size,model_file,params_file)                           #推理模型推理、预测
+    print(t)
+    
     
     
 if __name__ == "__main__":