[Model] Modify SR

examples/vision/sr/basicvsr/cpp/README.md

@@ -18,7 +18,7 @@ cmake .. -DFASTDEPLOY_INSTALL_DIR=${PWD}/fastdeploy-linux-x64-x.x.x
 make -j
 
 # 下载BasicVSR模型文件和测试视频
-wget https://bj.bcebos.com/paddlehub/fastdeploy/BasicVSR_reds_x4.tar
+wget https://bj.bcebos.com/paddlehub/fastdeploy/BasicVSR_reds_x4.tgz
 tar -xvf BasicVSR_reds_x4.tar
 wget https://bj.bcebos.com/paddlehub/fastdeploy/vsr_src.mp4
 

examples/vision/sr/basicvsr/cpp/infer.cc

@@ -20,8 +20,8 @@ const char sep = '\\';
 const char sep = '/';
 #endif
 
-void CpuInfer(const std::string& model_dir,
-              const std::string& video_file, int frame_num) {
+void CpuInfer(const std::string& model_dir, const std::string& video_file,
+              int frame_num) {
   auto model_file = model_dir + sep + "model.pdmodel";
   auto params_file = model_dir + sep + "model.pdiparams";
   auto model = fastdeploy::vision::sr::BasicVSR(model_file, params_file);
@@ -32,167 +32,165 @@ void CpuInfer(const std::string& model_dir,
   }
   // note: input/output shape is [b, n, c, h, w] (n = frame_nums; b=1(default))
   // b and n is dependent on export model shape
-  // see https://github.com/PaddlePaddle/PaddleGAN/blob/develop/docs/zh_CN/tutorials/video_super_resolution.md
+  // see
+  // https://github.com/PaddlePaddle/PaddleGAN/blob/develop/docs/zh_CN/tutorials/video_super_resolution.md
   cv::VideoCapture capture;
   // change your save video path
   std::string video_out_name = "output.mp4";
   capture.open(video_file);
-  if (!capture.isOpened())
-  {
-    std::cout<<"can not open video "<<std::endl;
+  if (!capture.isOpened()) {
+    std::cout << "can not open video " << std::endl;
     return;
   }
   // Get Video info :fps, frame count
   // it used 4.x version of opencv below
   // notice your opencv version and method of api.
   int video_fps = static_cast<int>(capture.get(cv::CAP_PROP_FPS));
-  int video_frame_count = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_COUNT));
+  int video_frame_count =
+      static_cast<int>(capture.get(cv::CAP_PROP_FRAME_COUNT));
   // Set fixed size for output frame, only for msvsr model
   int out_width = 1280;
   int out_height = 720;
-  std::cout << "fps: " << video_fps << "\tframe_count: " << video_frame_count << std::endl;
+  std::cout << "fps: " << video_fps << "\tframe_count: " << video_frame_count
+            << std::endl;
 
   // Create VideoWriter for output
   cv::VideoWriter video_out;
   std::string video_out_path("./");
   video_out_path += video_out_name;
   int fcc = cv::VideoWriter::fourcc('m', 'p', '4', 'v');
-  video_out.open(video_out_path, fcc, video_fps, cv::Size(out_width, out_height), true);
-  if (!video_out.isOpened())
-  {
+  video_out.open(video_out_path, fcc, video_fps,
+                 cv::Size(out_width, out_height), true);
+  if (!video_out.isOpened()) {
     std::cout << "create video writer failed!" << std::endl;
     return;
   }
   // Capture all frames and do inference
   cv::Mat frame;
   int frame_id = 0;
   bool reach_end = false;
-  while (capture.isOpened())
-  {
+  while (capture.isOpened()) {
     std::vector<cv::Mat> imgs;
-    for (int i = 0; i < frame_num; i++)
-    {
+    for (int i = 0; i < frame_num; i++) {
       capture.read(frame);
-      if (!frame.empty())
-      {
+      if (!frame.empty()) {
         imgs.push_back(frame);
-      }else{
+      } else {
         reach_end = true;
       }
     }
-    if (reach_end)
-    {
+    if (reach_end) {
       break;
     }
     std::vector<cv::Mat> results;
     model.Predict(imgs, results);
-    for (auto &item : results)
-    {
+    for (auto& item : results) {
       // cv::imshow("13",item);
       // cv::waitKey(30);
       video_out.write(item);
-      std::cout << "Processing frame: "<< frame_id << std::endl;
+      std::cout << "Processing frame: " << frame_id << std::endl;
       frame_id += 1;
     }
   }
-  std::cout << "inference finished, output video saved at " << video_out_path << std::endl;
+  std::cout << "inference finished, output video saved at " << video_out_path
+            << std::endl;
   capture.release();
   video_out.release();
 }
 
-void GpuInfer(const std::string& model_dir,
-              const std::string& video_file, int frame_num) {
+void GpuInfer(const std::string& model_dir, const std::string& video_file,
+              int frame_num) {
   auto model_file = model_dir + sep + "model.pdmodel";
   auto params_file = model_dir + sep + "model.pdiparams";
 
   auto option = fastdeploy::RuntimeOption();
   option.UseGpu();
-  auto model = fastdeploy::vision::sr::BasicVSR(
-      model_file, params_file, option);
+  auto model =
+      fastdeploy::vision::sr::BasicVSR(model_file, params_file, option);
 
   if (!model.Initialized()) {
     std::cerr << "Failed to initialize." << std::endl;
     return;
   }
   // note: input/output shape is [b, n, c, h, w] (n = frame_nums; b=1(default))
   // b and n is dependent on export model shape
-  // see https://github.com/PaddlePaddle/PaddleGAN/blob/develop/docs/zh_CN/tutorials/video_super_resolution.md
+  // see
+  // https://github.com/PaddlePaddle/PaddleGAN/blob/develop/docs/zh_CN/tutorials/video_super_resolution.md
   cv::VideoCapture capture;
   // change your save video path
   std::string video_out_name = "output.mp4";
   capture.open(video_file);
-  if (!capture.isOpened())
-  {
-    std::cout<<"can not open video "<<std::endl;
+  if (!capture.isOpened()) {
+    std::cout << "can not open video " << std::endl;
     return;
   }
   // Get Video info :fps, frame count
   int video_fps = static_cast<int>(capture.get(cv::CAP_PROP_FPS));
-  int video_frame_count = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_COUNT));
+  int video_frame_count =
+      static_cast<int>(capture.get(cv::CAP_PROP_FRAME_COUNT));
   // Set fixed size for output frame, only for msvsr model
   int out_width = 1280;
   int out_height = 720;
-  std::cout << "fps: " << video_fps << "\tframe_count: " << video_frame_count << std::endl;
+  std::cout << "fps: " << video_fps << "\tframe_count: " << video_frame_count
+            << std::endl;
 
   // Create VideoWriter for output
   cv::VideoWriter video_out;
   std::string video_out_path("./");
   video_out_path += video_out_name;
   int fcc = cv::VideoWriter::fourcc('m', 'p', '4', 'v');
-  video_out.open(video_out_path, fcc, video_fps, cv::Size(out_width, out_height), true);
-  if (!video_out.isOpened())
-  {
+  video_out.open(video_out_path, fcc, video_fps,
+                 cv::Size(out_width, out_height), true);
+  if (!video_out.isOpened()) {
     std::cout << "create video writer failed!" << std::endl;
     return;
   }
   // Capture all frames and do inference
   cv::Mat frame;
   int frame_id = 0;
   bool reach_end = false;
-  while (capture.isOpened())
-  {
+  while (capture.isOpened()) {
     std::vector<cv::Mat> imgs;
-    for (int i = 0; i < frame_num; i++)
-    {
+    for (int i = 0; i < frame_num; i++) {
       capture.read(frame);
-      if (!frame.empty())
-      {
+      if (!frame.empty()) {
         imgs.push_back(frame);
-      }else{
+      } else {
         reach_end = true;
       }
     }
-    if (reach_end)
-    {
+    if (reach_end) {
       break;
     }
     std::vector<cv::Mat> results;
     model.Predict(imgs, results);
-    for (auto &item : results)
-    {
+    for (auto& item : results) {
       // cv::imshow("13",item);
       // cv::waitKey(30);
       video_out.write(item);
-      std::cout << "Processing frame: "<< frame_id << std::endl;
+      std::cout << "Processing frame: " << frame_id << std::endl;
       frame_id += 1;
     }
   }
-  std::cout << "inference finished, output video saved at " << video_out_path << std::endl;
+  std::cout << "inference finished, output video saved at " << video_out_path
+            << std::endl;
   capture.release();
   video_out.release();
 }
 
-void TrtInfer(const std::string& model_dir,
-              const std::string& video_file, int frame_num) {
+void TrtInfer(const std::string& model_dir, const std::string& video_file,
+              int frame_num) {
   auto model_file = model_dir + sep + "model.pdmodel";
   auto params_file = model_dir + sep + "model.pdiparams";
   auto option = fastdeploy::RuntimeOption();
   option.UseGpu();
-  option.UseTrtBackend();
   // use paddle-TRT
+  option.UseTrtBackend();
+  option.EnablePaddleTrtCollectShape();
+  option.SetTrtInputShape("lrs", {1, 2, 3, 180, 320});
   option.EnablePaddleToTrt();
-  auto model = fastdeploy::vision::sr::BasicVSR(
-      model_file, params_file, option);
+  auto model =
+      fastdeploy::vision::sr::BasicVSR(model_file, params_file, option);
 
   if (!model.Initialized()) {
     std::cerr << "Failed to initialize." << std::endl;
@@ -201,81 +199,80 @@ void TrtInfer(const std::string& model_dir,
 
   // note: input/output shape is [b, n, c, h, w] (n = frame_nums; b=1(default))
   // b and n is dependent on export model shape
-  // see https://github.com/PaddlePaddle/PaddleGAN/blob/develop/docs/zh_CN/tutorials/video_super_resolution.md
+  // see
+  // https://github.com/PaddlePaddle/PaddleGAN/blob/develop/docs/zh_CN/tutorials/video_super_resolution.md
   cv::VideoCapture capture;
   // change your save video path
   std::string video_out_name = "output.mp4";
   capture.open(video_file);
-  if (!capture.isOpened())
-  {
-    std::cout<<"can not open video "<<std::endl;
+  if (!capture.isOpened()) {
+    std::cout << "can not open video " << std::endl;
     return;
   }
   // Get Video info :fps, frame count
   int video_fps = static_cast<int>(capture.get(cv::CAP_PROP_FPS));
-  int video_frame_count = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_COUNT));
+  int video_frame_count =
+      static_cast<int>(capture.get(cv::CAP_PROP_FRAME_COUNT));
   // Set fixed size for output frame, only for msvsr model
-  //Note that the resolution between the size and the original input is consistent when the model is exported,
+  // Note that the resolution between the size and the original input is
+  // consistent when the model is exported,
   // for example: [1,2,3,180,320], after 4x super separation [1,2,3,720,1080].
-  //Therefore, it is very important to derive the model
+  // Therefore, it is very important to derive the model
   int out_width = 1280;
   int out_height = 720;
-  std::cout << "fps: " << video_fps << "\tframe_count: " << video_frame_count << std::endl;
+  std::cout << "fps: " << video_fps << "\tframe_count: " << video_frame_count
+            << std::endl;
 
   // Create VideoWriter for output
   cv::VideoWriter video_out;
   std::string video_out_path("./");
   video_out_path += video_out_name;
   int fcc = cv::VideoWriter::fourcc('m', 'p', '4', 'v');
-  video_out.open(video_out_path, fcc, video_fps, cv::Size(out_width, out_height), true);
-  if (!video_out.isOpened())
-  {
+  video_out.open(video_out_path, fcc, video_fps,
+                 cv::Size(out_width, out_height), true);
+  if (!video_out.isOpened()) {
     std::cout << "create video writer failed!" << std::endl;
     return;
   }
   // Capture all frames and do inference
   cv::Mat frame;
   int frame_id = 0;
   bool reach_end = false;
-  while (capture.isOpened())
-  {
+  while (capture.isOpened()) {
     std::vector<cv::Mat> imgs;
-    for (int i = 0; i < frame_num; i++)
-    {
+    for (int i = 0; i < frame_num; i++) {
       capture.read(frame);
-      if (!frame.empty())
-      {
+      if (!frame.empty()) {
         imgs.push_back(frame);
-      }else{
+      } else {
         reach_end = true;
       }
     }
-    if (reach_end)
-    {
+    if (reach_end) {
       break;
     }
     std::vector<cv::Mat> results;
     model.Predict(imgs, results);
-    for (auto &item : results)
-    {
+    for (auto& item : results) {
       // cv::imshow("13",item);
       // cv::waitKey(30);
       video_out.write(item);
-      std::cout << "Processing frame: "<< frame_id << std::endl;
+      std::cout << "Processing frame: " << frame_id << std::endl;
       frame_id += 1;
     }
   }
-  std::cout << "inference finished, output video saved at " << video_out_path << std::endl;
+  std::cout << "inference finished, output video saved at " << video_out_path
+            << std::endl;
   capture.release();
   video_out.release();
 }
 
 int main(int argc, char* argv[]) {
   if (argc < 4) {
-    std::cout
-        << "Usage: infer_demo path/to/model_dir path/to/video frame number run_option, "
-           "e.g ./infer_model ./vsr_model_dir ./person.mp4 0 2"
-        << std::endl;
+    std::cout << "Usage: infer_demo path/to/model_dir path/to/video frame "
+                 "number run_option, "
+                 "e.g ./infer_model ./vsr_model_dir ./vsr_src.mp4 0 2"
+              << std::endl;
     std::cout << "The data type of run_option is int, 0: run with cpu; 1: run "
                  "with gpu; 2: run with gpu and use tensorrt backend."
               << std::endl;

examples/vision/sr/basicvsr/python/README.md

@@ -13,15 +13,15 @@ git clone https://github.com/PaddlePaddle/FastDeploy.git
 cd FastDeploy/examples/vision/sr/basicvsr/python
 
 # 下载BasicVSR模型文件和测试视频
-wget https://bj.bcebos.com/paddlehub/fastdeploy/BasicVSR_reds_x4.tar
+wget https://bj.bcebos.com/paddlehub/fastdeploy/BasicVSR_reds_x4.tgz
 tar -xvf BasicVSR_reds_x4.tar
 wget https://bj.bcebos.com/paddlehub/fastdeploy/vsr_src.mp4
 # CPU推理
-python infer.py --model BasicVSR_reds_x4 --video person.mp4 --frame_num 2 --device cpu
+python infer.py --model BasicVSR_reds_x4 --video vsr_src.mp4 --frame_num 2 --device cpu
 # GPU推理
-python infer.py --model BasicVSR_reds_x4 --video person.mp4 --frame_num 2 --device gpu
+python infer.py --model BasicVSR_reds_x4 --video vsr_src.mp4 --frame_num 2 --device gpu
 # GPU上使用TensorRT推理 （注意：TensorRT推理第一次运行，有序列化模型的操作，有一定耗时，需要耐心等待）
-python infer.py --model BasicVSR_reds_x4 --video person.mp4 --frame_num 2 --device gpu --use_trt True
+python infer.py --model BasicVSR_reds_x4 --video vsr_src.mp4 --frame_num 2 --device gpu --use_trt True
 ```
 
 ## BasicVSR Python接口

examples/vision/sr/basicvsr/python/infer.py

@@ -30,6 +30,8 @@ def build_option(args):
         option.use_gpu()
     if args.use_trt:
         option.use_trt_backend()
+        option.enable_paddle_trt_collect_shape()
+        option.set_trt_input_shape("lrs", [1, 2, 3, 180, 320])
         option.enable_paddle_to_trt()
     return option
 
@@ -56,7 +58,7 @@ def build_option(args):
 # Create VideoWriter for output
 video_out_dir = "./"
 video_out_path = os.path.join(video_out_dir, video_out_name)
-fucc = cv2.VideoWriter_fourcc(*"mp4v")
+fucc = cv2.VideoWriter_fourcc(* "mp4v")
 video_out = cv2.VideoWriter(video_out_path, fucc, video_fps,
                             (out_width, out_height), True)
 if not video_out.isOpened():