Add files via upload

Author: fireindark707

cal_column_similarity.py

@@ -9,24 +9,39 @@
 parser.add_argument("-p","--path", help="path to the folder containing the test data")
 parser.add_argument("-m", "--model", help="path to the model")
 parser.add_argument("-t", "--threshold", help="threshold for inference")
+parser.add_argument("-s", "--strategy", help="one-to-one or one-to-many", default="one-to-many")
 args = parser.parse_args()
 
-def create_similarity_matrix(pth,pred):
+def create_similarity_matrix(pth,preds,pred_labels_list,strategy="one-to-many"):
     """
     Create a similarity matrix from the prediction
     """
+    preds = np.array(preds)
+    preds = np.mean(preds,axis=0)
+    pred_labels_list = np.array(pred_labels_list)
+    pred_labels = np.mean(pred_labels_list,axis=0)
+    pred_labels = np.where(pred_labels>0.5,1,0)
+    # read column names
     df1 = pd.read_csv(pth+"/Table1.csv")
     df2 = pd.read_csv(pth+"/Table2.csv")
     df1_cols = df1.columns
     df2_cols = df2.columns
-    # create similarity matrix for pred values
-    sim_matrix = np.zeros((len(df1_cols),len(df2_cols)))
-    for i in range(len(df1_cols)):
-        for j in range(len(df2_cols)):
-            sim_matrix[i,j] = pred[i*len(df2_cols)+j]
-    # create dataframe
-    df = pd.DataFrame(sim_matrix,index=df1_cols,columns=df2_cols)
-    return df
+    # create similarity matrix for pred values 
+    preds_matrix = np.array(preds).reshape(len(df1_cols),len(df2_cols))
+    if strategy == "one-to-many":
+        pred_labels_matrix = np.array(pred_labels).reshape(len(df1_cols),len(df2_cols))
+    elif strategy == "one-to-one":
+        pred_labels_matrix = np.zeros((len(df1_cols),len(df2_cols)))
+        for i in range(len(df1_cols)):
+            for j in range(len(df2_cols)):
+                if pred_labels[i*len(df2_cols)+j] == 1:
+                    max_row = max(preds_matrix[i,:])
+                    max_col = max(preds_matrix[:,j])
+                    if preds_matrix[i,j] == max_row and preds_matrix[i,j] == max_col:
+                        pred_labels_matrix[i,j] = 1
+    df_pred = pd.DataFrame(preds_matrix,columns=df2_cols,index=df1_cols)
+    df_pred_labels = pd.DataFrame(pred_labels_matrix,columns=df2_cols,index=df1_cols)
+    return df_pred,df_pred_labels
 
 if __name__ == '__main__':
     pth = args.path
@@ -48,13 +63,7 @@ def create_similarity_matrix(pth,pred):
         preds.append(pred)
         pred_labels_list.append(pred_labels)
         del bst
-    preds = np.array(preds)
-    preds = np.mean(preds,axis=0)
-    pred_labels_list = np.array(pred_labels_list)
-    pred_labels = np.mean(pred_labels_list,axis=0)
-    pred_labels = np.where(pred_labels>0.5,1,0)
 
-    df_pred = create_similarity_matrix(pth,pred)
-    df_pred_labels = create_similarity_matrix(pth,pred_labels)
-    df_pred.to_csv(pth+"/similarity_matrix_value.csv")
-    df_pred_labels.to_csv(pth+"/similarity_matrix_label.csv")
+    df_pred,df_pred_labels = create_similarity_matrix(pth,preds,pred_labels_list,strategy=args.strategy)
+    df_pred.to_csv(pth+"/similarity_matrix_value.csv",index=True)
+    df_pred_labels.to_csv(pth+"/similarity_matrix_label.csv",index=True)