multiprocessing test

Author: argusswift

config/yolov4_config.py

@@ -7,7 +7,7 @@
 
 
 MODEL_TYPE = {
-    "TYPE": "Mobilenetv3-YOLOv4"
+    "TYPE": "Mobilenet-YOLOv4"
 }  # YOLO type:YOLOv4, Mobilenet-YOLOv4 or Mobilenetv3-YOLOv4
 
 CONV_TYPE = {"TYPE": "DO_CONV"}  # conv type:DO_CONV or GENERAL
@@ -16,7 +16,7 @@
 
 # train
 TRAIN = {
-    "DATA_TYPE": "Customer",  # DATA_TYPE: VOC ,COCO or Customer
+    "DATA_TYPE": "VOC",  # DATA_TYPE: VOC ,COCO or Customer
     "TRAIN_IMG_SIZE": 416,
     "AUGMENT": True,
     "BATCH_SIZE": 1,

eval/evaluator.py

@@ -7,7 +7,9 @@
 from utils.heatmap import imshowAtt
 import config.yolov4_config as cfg
 import time
-
+import multiprocessing
+from multiprocessing.dummy import Pool as ThreadPool        # 线程池
+from collections import defaultdict
 current_milli_time = lambda: int(round(time.time() * 1000))
 
 
@@ -28,11 +30,14 @@ def __init__(self, model=None, showatt=False):
         self.val_shape = cfg.VAL["TEST_IMG_SIZE"]
         self.model = model
         self.device = next(model.parameters()).device
-        self.__visual_imgs = 0
+        self.visual_imgs = 0
+        self.multi_scale_test = cfg.VAL["MULTI_SCALE_VAL"]
+        self.flip_test = cfg.VAL["FLIP_VAL"]
         self.showatt = showatt
         self.inference_time = 0.0
+        self.final_result = defaultdict(list)
 
-    def APs_voc(self, multi_test=False, flip_test=False):
+    def APs_voc(self):
         img_inds_file = os.path.join(
             self.val_data_path, "ImageSets", "Main", "test.txt"
         )
@@ -47,47 +52,47 @@ def APs_voc(self, multi_test=False, flip_test=False):
         if not os.path.exists(output_path):
             os.mkdir(output_path)
         os.mkdir(self.pred_result_path)
-        for img_ind in tqdm(img_inds):
-            img_path = os.path.join(
-                self.val_data_path, "JPEGImages", img_ind + ".jpg"
-            )
-            img = cv2.imread(img_path)
-            bboxes_prd = self.get_bbox(img, multi_test, flip_test)
-
-            f = open("./output/" + img_ind + ".txt", "w")
-            for bbox in bboxes_prd:
-                coor = np.array(bbox[:4], dtype=np.int32)
-                score = bbox[4]
-                class_ind = int(bbox[5])
-
-                class_name = self.classes[class_ind]
-                score = "%.4f" % score
-                xmin, ymin, xmax, ymax = map(str, coor)
-                s = " ".join([img_ind, score, xmin, ymin, xmax, ymax]) + "\n"
-
-                with open(
-                    os.path.join(
-                        self.pred_result_path,
-                        "comp4_det_test_" + class_name + ".txt",
-                    ),
-                    "a",
-                ) as r:
-                    r.write(s)
-                f.write(
-                    "%s %s %s %s %s %s\n"
-                    % (
-                        class_name,
-                        score,
-                        str(xmin),
-                        str(ymin),
-                        str(xmax),
-                        str(ymax),
-                    )
-                )
-            f.close()
+        imgs_count = len(img_inds)
+        cpu_nums = multiprocessing.cpu_count()
+        pool = ThreadPool(cpu_nums)
+        with tqdm(total=imgs_count) as pbar:
+            for i, _ in enumerate(pool.imap_unordered(self.Single_APs_voc, img_inds)):
+                pbar.update()
+        for class_name in self.final_result:
+            with open(os.path.join(self.pred_result_path, 'comp4_det_test_' + class_name + '.txt'), 'a') as f:
+                str_result = ''.join(self.final_result[class_name])
+                f.write(str_result)
         self.inference_time = 1.0 * self.inference_time / len(img_inds)
         return self.__calc_APs(), self.inference_time
 
+    def Single_APs_voc(self, img_ind):
+        img_path = os.path.join(self.val_data_path, 'JPEGImages', img_ind + '.jpg')
+        img = cv2.imread(img_path)
+        bboxes_prd = self.get_bbox(img, self.multi_scale_test, self.flip_test)
+
+        if bboxes_prd.shape[0] != 0  and self.visual_imgs < 100:
+            boxes = bboxes_prd[..., :4]
+            class_inds = bboxes_prd[..., 5].astype(np.int32)
+            scores = bboxes_prd[..., 4]
+
+            visualize_boxes(image=img, boxes=boxes, labels=class_inds, probs=scores, class_labels=self.classes)
+            path = os.path.join(cfg.PROJECT_PATH, "data/results/{}.jpg".format(self.visual_imgs))
+            cv2.imwrite(path, img)
+
+            self.visual_imgs += 1
+
+        for bbox in bboxes_prd:
+            coor = np.array(bbox[:4], dtype=np.int32)
+            score = bbox[4]
+            class_ind = int(bbox[5])
+
+            class_name = self.classes[class_ind]
+            score = '%.4f' % score
+            xmin, ymin, xmax, ymax = map(str, coor)
+            result = ' '.join([img_ind, score, xmin, ymin, xmax, ymax]) + '\n'
+
+            self.final_result[class_name].append(result)
+
     def get_bbox(self, img, multi_test=False, flip_test=False, mode=None):
         if multi_test:
             test_input_sizes = range(320, 640, 96)

eval_voc.py

@@ -25,8 +25,6 @@ def __init__(
         self.__conf_threshold = cfg.VAL["CONF_THRESH"]
         self.__nms_threshold = cfg.VAL["NMS_THRESH"]
         self.__device = gpu.select_device(gpu_id)
-        self.__multi_scale_val = cfg.VAL["MULTI_SCALE_VAL"]
-        self.__flip_val = cfg.VAL["FLIP_VAL"]
         self.__showatt = showatt
         self.__visiual = visiual
         self.__eval = eval
@@ -57,7 +55,7 @@ def val(self):
             with torch.no_grad():
                 APs, inference_time = Evaluator(
                     self.__model, showatt=False
-                ).APs_voc(self.__multi_scale_val, self.__flip_val)
+                ).APs_voc()
                 for i in APs:
                     logger.info("{} --> mAP : {}".format(i, APs[i]))
                     mAP += APs[i]

train.py

@@ -318,7 +318,7 @@ def train(self):
     parser.add_argument(
         "--weight_path",
         type=str,
-        default="weight/mobilenetv3.pth",
+        default="weight/mobilenetv2.pth",
         help="weight file path",
     )  # weight/darknet53_448.weights
     parser.add_argument(
@@ -330,7 +330,7 @@ def train(self):
     parser.add_argument(
         "--gpu_id",
         type=int,
-        default=-1,
+        default=0,
         help="whither use GPU(0) or CPU(-1)",
     )
     parser.add_argument("--log_path", type=str, default="log/", help="log path")
@@ -349,7 +349,7 @@ def train(self):
     parser.add_argument(
         "--showatt",
         type=bool,
-        default=True,
+        default=False,
         help="whether to show attention map"
     )
     opt = parser.parse_args()