update

rov_train.py

@@ -137,11 +137,12 @@ def pack_result_to_csv(filename, sort_columns=None, filepath=config_.DATA_DIR_PA
     df.to_csv(file, index=False)
 
 
-def pack_list_result_to_csv(filename, data, sort_columns=None, filepath=config_.DATA_DIR_PATH, ascending=True):
+def pack_list_result_to_csv(filename, data, columns=None, sort_columns=None, filepath=config_.DATA_DIR_PATH, ascending=True):
     """
     打包数据并存入csv, 数据为字典列表
     :param filename: csv文件名
     :param data: 数据，type-list [{}, {},...]
+    :param columns: 列名顺序
     :param sort_columns: 指定排序列名列名，type-list, 默认为None
     :param filepath: csv文件存放路径，默认为config_.DATA_DIR_PATH
     :param ascending: 是否按指定列的数组升序排列，默认为True，即升序排列
@@ -153,7 +154,7 @@ def pack_list_result_to_csv(filename, data, sort_columns=None, filepath=config_.
     df = pd.DataFrame(data=data)
     if sort_columns:
         df = df.sort_values(by=sort_columns, ascending=ascending)
-    df.to_csv(file, index=False)
+    df.to_csv(file, index=False, columns=columns)
 
 
 def predict():
@@ -192,8 +193,11 @@ def predict():
 
     # 打包预测结果存入csv
     predict_result_filename = 'predict.csv'
-    pack_list_result_to_csv(filename=predict_result_filename, sort_columns=['rov_score'],
-                            ascending=False, data=predict_result)
+    pack_list_result_to_csv(filename=predict_result_filename,
+                            data=predict_result,
+                            columns=['video_id', 'rov_score', 'normal_y_', 'y_', 'y'],
+                            sort_columns=['rov_score'],
+                            ascending=False)
 
     # 上传redis
     key_name = config_.RECALL_KEY_NAME_PREFIX + time.strftime('%Y%m%d')