first

fine/util/test_tool.py

@@ -0,0 +1,99 @@
+import os
+import math
+import numpy as np
+import cv2
+import datetime
+
+import torch
+from torchvision.utils import save_image, make_grid
+
+import pdb
+
+def dt():
+    return datetime.datetime.now().strftime("%Y-%m-%d-%H_%M_%S")
+
+def set_text(text, img):
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    if isinstance(text, str):
+        cont = text
+        cv2.putText(img, cont, (20, 50), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+    if isinstance(text, float):
+        cont = '%.4f'%text
+        cv2.putText(img, cont, (20, 50), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+    if isinstance(text, list):
+        for count in range(len(img)):
+            cv2.putText(img[count], text[count], (20, 50), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+    return img
+
+def save_multi_img(img_list, text_list, grid_size=[5,5], sub_size=200, save_dir='./', save_name=None):
+    if len(img_list) > grid_size[0]*grid_size[1]:
+        merge_height = math.ceil(len(img_list) / grid_size[0]) * sub_size
+    else:
+        merge_height = grid_size[1]*sub_size
+    merged_img = np.zeros((merge_height, grid_size[0]*sub_size, 3))
+
+    if isinstance(img_list[0], str):
+        img_name_list = img_list
+        img_list = []
+        for img_name in img_name_list:
+            img_list.append(cv2.imread(img_name))
+
+    img_counter = 0
+    for img, txt in zip(img_list, text_list):
+        img = cv2.resize(img, (sub_size, sub_size))
+        img = set_text(txt, img)
+        pos = [img_counter // grid_size[1], img_counter % grid_size[1]]
+        sub_pos = [pos[0]*sub_size, (pos[0]+1)*sub_size,
+                   pos[1]*sub_size, (pos[1]+1)*sub_size]
+        merged_img[sub_pos[0]:sub_pos[1], sub_pos[2]:sub_pos[3], :] = img
+        img_counter += 1
+
+    if save_name is None:
+        img_save_path = os.path.join(save_dir, dt()+'.png')
+    else:
+        img_save_path = os.path.join(save_dir, save_name+'.png')
+    cv2.imwrite(img_save_path,  merged_img)
+    print('saved img in %s ...'%img_save_path)
+
+
+def cls_base_acc(result_gather):
+    top1_acc = {}
+    top3_acc = {}
+    cls_count = {}
+    for img_item in result_gather.keys():
+        acc_case = result_gather[img_item]
+
+        if acc_case['label'] in cls_count:
+            cls_count[acc_case['label']] += 1
+            if acc_case['top1_cat'] == acc_case['label']:
+                top1_acc[acc_case['label']] += 1
+            if acc_case['label'] in [acc_case['top1_cat'], acc_case['top2_cat'], acc_case['top3_cat']]:
+                top3_acc[acc_case['label']] += 1
+        else:
+            cls_count[acc_case['label']] = 1
+            if acc_case['top1_cat'] == acc_case['label']:
+                top1_acc[acc_case['label']] = 1
+            else:
+                top1_acc[acc_case['label']] = 0
+
+            if acc_case['label'] in [acc_case['top1_cat'], acc_case['top2_cat'], acc_case['top3_cat']]:
+                top3_acc[acc_case['label']] = 1
+            else:
+                top3_acc[acc_case['label']] = 0
+
+    for label_item in cls_count:
+        top1_acc[label_item] /= max(1.0*cls_count[label_item], 0.001)
+        top3_acc[label_item] /= max(1.0*cls_count[label_item], 0.001)
+
+    print('top1_acc:', top1_acc)
+    print('top3_acc:', top3_acc)
+    print('cls_count', cls_count)
+
+    return top1_acc, top3_acc, cls_count
+
+
+
+
+
+
+