feat(map): voc map计算

Author: zjZSTU

py/batch_detect.py

@@ -9,17 +9,14 @@
 
 import os
 import glob
-import time
-import shutil
 import cv2
 import numpy as np
 import torch
 import torchvision.transforms as transforms
 
 from utils import file
 from utils import util
-from utils import draw
-from models.location_dataset import LocationDataset
+from utils import voc_map
 from models.yolo_v1 import YOLO_v1
 
 S = 7
@@ -28,6 +25,18 @@
 
 cate_list = ['cucumber', 'eggplant', 'mushroom']
 
+dst_root_dir = './data/outputs'
+dst_target_dir = os.path.join(dst_root_dir, 'targets')
+dst_pred_dir = os.path.join(dst_root_dir, 'preds')
+dst_img_dir = os.path.join(dst_root_dir, 'imgs')
+tmp_json_dir = os.path.join(dst_root_dir, '.tmp_files')
+
+file.check_dir(dst_root_dir)
+file.check_dir(dst_target_dir)
+file.check_dir(dst_pred_dir)
+file.check_dir(dst_img_dir)
+file.check_dir(tmp_json_dir)
+
 
 def get_transform():
     transform = transforms.Compose([
@@ -41,9 +50,11 @@ def get_transform():
 
 
 def load_data(root_dir):
-    img_path_list = glob.glob(os.path.join(root_dir, '*.jpg'))
-    annotation_path_list = [os.path.join(root_dir, os.path.splitext(os.path.basename(img_path))[0] + ".xml")
-                            for img_path in img_path_list]
+    img_path_list = glob.glob(os.path.join(root_dir, 'imgs', '*.jpg'))
+    img_path_list.sort()
+    annotation_path_list = [
+        os.path.join(root_dir, 'annotations', os.path.splitext(os.path.basename(img_path))[0] + ".xml")
+        for img_path in img_path_list]
 
     return img_path_list, annotation_path_list
 
@@ -137,16 +148,6 @@ def save_data(img_name, img, target_cates, target_bboxs, pred_cates, pred_probs,
     :param pred_probs: 预测边界框置信度
     :param pred_bboxs: 预测边界框坐标
     """
-    dst_root_dir = './data/outputs'
-    dst_target_dir = os.path.join(dst_root_dir, 'targets')
-    dst_pred_dir = os.path.join(dst_root_dir, 'preds')
-    dst_img_dir = os.path.join(dst_root_dir, 'imgs')
-
-    file.check_dir(dst_root_dir)
-    file.check_dir(dst_target_dir)
-    file.check_dir(dst_pred_dir)
-    file.check_dir(dst_img_dir)
-
     img_path = os.path.join(dst_img_dir, img_name + ".png")
     cv2.imwrite(img_path, img)
     annotation_path = os.path.join(dst_target_dir, img_name + ".txt")
@@ -176,7 +177,7 @@ def save_data(img_name, img, target_cates, target_bboxs, pred_cates, pred_probs,
     model = load_model(device)
 
     transform = get_transform()
-    img_path_list, annotation_path_list = load_data('./data/training_images')
+    img_path_list, annotation_path_list = load_data('./data/location_dataset')
     # print(img_path_list)
 
     N = len(img_path_list)
@@ -217,4 +218,5 @@ def save_data(img_name, img, target_cates, target_bboxs, pred_cates, pred_probs,
         img_name = os.path.splitext(os.path.basename(img_path))[0]
         save_data(img_name, data_dict['src'], data_dict['name_list'], data_dict['bndboxs'],
                   pred_cates, pred_cate_probs, pred_bboxs)
-    print('done')
+    print('compute mAP')
+    voc_map.voc_map(dst_target_dir, dst_pred_dir, tmp_json_dir)

py/lib/utils/draw.py

@@ -59,7 +59,7 @@ def plot_bboxs(img, bndboxs, name_list, pred_boxs, pred_cates, pred_probs):
 
         xmin, ymin, xmax, ymax = np.array(bbox, dtype=np.int)
         cv2.rectangle(dst, (xmin, ymin), (xmax, ymax), (0, 0, 255), thickness=1)
-        cv2.putText(dst, '%s_%f' % (cate_list[cate], prob), (xmin, ymax), 1, cv2.FONT_HERSHEY_PLAIN, (0, 0, 255),
+        cv2.putText(dst, '%s_%.3f' % (cate_list[cate], prob), (xmin, ymin), 1, cv2.FONT_HERSHEY_PLAIN, (0, 0, 255),
                     thickness=1)
 
     return dst

py/lib/utils/file.py

@@ -11,11 +11,15 @@
 import xmltodict
 import numpy as np
 import torch
+import shutil
+import json
+import glob
 
 
 def check_dir(data_dir):
-    if not os.path.exists(data_dir):
-        os.mkdir(data_dir)
+    if os.path.exists(data_dir):
+        shutil.rmtree(data_dir)
+    os.mkdir(data_dir)
 
 
 def parse_location_xml(xml_path):
@@ -72,6 +76,79 @@ def parse_location_xml(xml_path):
         return np.array(bndboxs), name_list
 
 
+def file_lines_to_list(path):
+    """
+     Convert the lines of a file to a list
+    """
+    # open txt file lines to a list
+    with open(path) as f:
+        content = f.readlines()
+    # remove whitespace characters like `\n` at the end of each line
+    content = [x.strip() for x in content]
+    return content
+
+
+def parse_ground_truth(ground_truth_dir, tmp_json_dir):
+    """
+    解析每个图片的真值边界框，以格式{"cate": "cucumber", "bbox": [23, 42, 206, 199], "used": true}保存
+    """
+    gt_path_list = glob.glob(os.path.join(ground_truth_dir, '*.txt'))
+
+    # 统计每类的真值标注框数量
+    gt_per_classes_dict = {}
+    for gt_path in gt_path_list:
+        json_list = list()
+        lines = file_lines_to_list(gt_path)
+        for line in lines:
+            cate, xmin, ymin, xmax, ymax = line.split(' ')
+            json_list.append({'cate': cate, 'bbox': [int(xmin), int(ymin), int(xmax), int(ymax)], 'used': False})
+
+            if gt_per_classes_dict.get(cate) is None:
+                gt_per_classes_dict[cate] = 1
+            else:
+                gt_per_classes_dict[cate] += 1
+        # 保存
+        name = os.path.splitext(os.path.basename(gt_path))[0]
+        json_path = os.path.join(tmp_json_dir, name + ".json")
+        with open(json_path, 'w') as f:
+            json.dump(json_list, f)
+
+    return gt_per_classes_dict
+
+
+def parse_detection_results(detection_result_dir, tmp_json_dir):
+    """
+    解析每个类别的预测边界框，以格式{"confidence": "0.999", "file_id": "cucumber_61", "bbox": [16, 42, 225, 163]}保存
+    """
+    dr_path_list = glob.glob(os.path.join(detection_result_dir, '*.txt'))
+
+    # 保存每个类别的预测边界框信息
+    dt_per_classes_dict = dict()
+    for dr_path in dr_path_list:
+        lines = file_lines_to_list(dr_path)
+        name = os.path.splitext(os.path.basename(dr_path))[0]
+
+        for line in lines:
+            cate, confidence, xmin, ymin, xmax, ymax = line.split(' ')
+            if dt_per_classes_dict.get(cate) is None:
+                dt_per_classes_dict[cate] = [
+                    {'confidence': confidence, 'file_id': name, 'bbox': [int(xmin), int(ymin), int(xmax), int(ymax)]}]
+            else:
+                dt_per_classes_dict[cate].append(
+                    {'confidence': confidence, 'file_id': name, 'bbox': [int(xmin), int(ymin), int(xmax), int(ymax)]})
+
+    # 保存
+    for key, value in dt_per_classes_dict.items():
+        # 按置信度递减排序
+        value.sort(key=lambda x: float(x['confidence']), reverse=True)
+
+        json_path = os.path.join(tmp_json_dir, key + "_dt.json")
+        with open(json_path, 'w') as f:
+            json.dump(value, f)
+
+    return dt_per_classes_dict
+
+
 def save_model(model_weights, model_save_path):
     torch.save(model_weights, model_save_path)