mod draw2 with color

draw_confusion.py

@@ -3,23 +3,40 @@
 import seaborn as sns
 from sklearn.metrics import confusion_matrix
 import matplotlib.pyplot as plt
+from colorama import Fore,Back,Style
 
 
 def draw(y_true, y_pred):
+	min_len = min( len(y_true), len(y_pred) )
+	
 	sns.set()
 	f, ax=plt.subplots()
 	# y_true = [0,0,1,2,1,2,0,2,2,0,1,1]
 	# y_pred = [1,0,1,2,1,0,0,2,2,0,1,1]
-	cm = confusion_matrix(y_true, y_pred, labels=[-1] * 13)
+	cm = confusion_matrix(y_true, y_pred, labels=[-1] * min_len)
 	print(cm) #打印出来看看
 	sns.heatmap(cm, annot=True, ax=ax) #画热力图
 
 def draw2(y_true, y_pred):
 	min_len = min( len(y_true), len(y_pred) )
+	
+	print('\t', end='')
+	for i in range(min_len):
+		y_true_format = str(y_true[i])[:3]
+		print('%s\t' % y_true_format, end='')
+	print('')
+
 	for i in range(min_len):
+		print(Fore.RESET + '%s\t' % str(i + 1), end='')
 		for j in range(i):
 			print('\t', end='')
-		print(str(y_pred[i]))
 
-draw2( [1.0] * 13, [0.5] * 13 )
+		# print with color
+		if y_pred[i] > 0.5:
+			print(Fore.GREEN + str(y_pred[i]))
+		else:
+			print(Fore.RED + str(y_pred[i]))
+
+if __name__ == '__main__':
+	draw2( [1.0] * 13, [0.5] * 13 )
 

requirements.txt

@@ -0,0 +1 @@
+colorama

test_dreamtalk.py

@@ -1,43 +1,58 @@
 import requests
 import sys
 from draw_confusion import draw, draw2
+from tqdm import tqdm
 
+
+DEBUG = True #False
 API_URL = "https://api-inference.huggingface.co/models/MIT/ast-finetuned-audioset-10-10-0.4593"
 headers = {"Authorization": "Bearer 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"}
 
+
+# 处理请求
 # filename = '1.flac'
-def query(filename):
+def request_api(filename):
     with open(filename, "rb") as f:
         data = f.read()
     response = requests.post(API_URL, headers=headers, data=data)
     return response.json()
 
-# res = query(filename)
-# print(str(res))
-
-# 处理命令行
-if __name__ == "__main__":
-    # 获取命令行参数
-    if len(sys.argv) < 2:
-        print("用法：python x.py <文件或通配符>")
-        sys.exit(1)
 
-    y_len = 13 #len(sys.argv[1:]:)
+# 批量处理
+def batch_request_api(file_uris):
+    if DEBUG: print('batch_request_api')
+    y_len = len(file_uris)
     y_true = [0] * y_len
     y_pred = [0] * y_len
     y_idx = 0
-    for input_file in sys.argv[1:]:
-        res = query(input_file)
-        print('%s:' % str(input_file))
-        print('%s' % str(res[:3]))
+    for input_file in tqdm(file_uris):
+        res = request_api(input_file)
+        # print('%s %s:' % (str(y_idx), str(input_file)) )
+        # print('%s' % str(res[:3]))
         
         first_label = str(res[0]['label'])
         first_score = res[0]['score']
-        print(str(first_label))
-        print(str(first_score))
+        # print(str(first_label))
+        # print(str(first_score))
 
-        y_true[y_idx] = 1.0
+        y_true[y_idx] = first_label
         y_pred[y_idx] = round(first_score, 1)
         y_idx = y_idx + 1
     
-    draw2(y_true, y_pred)
+    return y_true, y_pred
+
+
+# 处理命令行
+if __name__ == "__main__":
+	if DEBUG: print('main, ' + str(sys.argv[1:]))	
+	if DEBUG: print('main, ' + str(len(sys.argv)))
+
+	# 获取命令行参数
+	if len(sys.argv) < 2:
+		print("用法：python x.py <文件或通配符>")
+		sys.exit(1)
+	
+	if DEBUG: print('main, batch_request_api')
+	y_true, y_pred = batch_request_api(sys.argv[1:])
+	if DEBUG: print('y_true = %s, y_pred = %s' % (str(y_true), str(y_pred)))
+	draw2(y_true, y_pred)