Merge pull request WZMIAOMIAO#177 from WZMIAOMIAO/dev

Dev
zyc961018 · Mar 13, 2021 · 40616eb · 40616eb
2 parents 8fbf8b8 + 1fe1407
commit 40616eb
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 5 deletions.
diff --git a/pytorch_object_detection/faster_rcnn/my_dataset.py b/pytorch_object_detection/faster_rcnn/my_dataset.py
@@ -49,17 +49,22 @@ def __getitem__(self, idx):
         image = Image.open(img_path)
         if image.format != "JPEG":
             raise ValueError("Image '{}' format not JPEG".format(img_path))
+
         boxes = []
         labels = []
         iscrowd = []
+        assert "object" in data, "{} lack of object information.".format(xml_path)
         for obj in data["object"]:
             xmin = float(obj["bndbox"]["xmin"])
             xmax = float(obj["bndbox"]["xmax"])
             ymin = float(obj["bndbox"]["ymin"])
             ymax = float(obj["bndbox"]["ymax"])
             boxes.append([xmin, ymin, xmax, ymax])
             labels.append(self.class_dict[obj["name"]])
-            iscrowd.append(int(obj["difficult"]))
+            if "difficult" in obj:
+                iscrowd.append(int(obj["difficult"]))
+            else:
+                iscrowd.append(0)
 
         # convert everything into a torch.Tensor
         boxes = torch.as_tensor(boxes, dtype=torch.float32)

diff --git a/pytorch_object_detection/ssd/my_dataset.py b/pytorch_object_detection/ssd/my_dataset.py
@@ -45,6 +45,8 @@ def __getitem__(self, idx):
         image = Image.open(img_path)
         if image.format != "JPEG":
             raise ValueError("Image '{}' format not JPEG".format(img_path))
+
+        assert "object" in data, "{} lack of object information.".format(xml_path)
         boxes = []
         labels = []
         iscrowd = []
@@ -56,7 +58,10 @@ def __getitem__(self, idx):
             ymax = float(obj["bndbox"]["ymax"]) / data_height
             boxes.append([xmin, ymin, xmax, ymax])
             labels.append(self.class_dict[obj["name"]])
-            iscrowd.append(int(obj["difficult"]))
+            if "difficult" in obj:
+                iscrowd.append(int(obj["difficult"]))
+            else:
+                iscrowd.append(0)
 
         # convert everything into a torch.Tensor
         boxes = torch.as_tensor(boxes, dtype=torch.float32)

diff --git a/pytorch_object_detection/yolov3_spp/build_utils/datasets.py b/pytorch_object_detection/yolov3_spp/build_utils/datasets.py
@@ -613,8 +613,8 @@ def letterbox(img: np.ndarray,
     new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
     dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
     if auto:  # minimun rectangle 保证原图比例不变，将图像最大边缩放到指定大小
-        # 这里的取余操作可以保证padding后的图片是32的整数倍(416x416)，如果是(512x512)可以保证是64的整数倍
-        dw, dh = np.mod(dw, 64), np.mod(dh, 64)  # wh padding
+        # 这里的取余操作可以保证padding后的图片是32的整数倍
+        dw, dh = np.mod(dw, 32), np.mod(dh, 32)  # wh padding
     elif scale_fill:  # stretch 简单粗暴的将图片缩放到指定尺寸
         dw, dh = 0, 0
         new_unpad = new_shape

diff --git a/pytorch_object_detection/yolov3_spp/build_utils/utils.py b/pytorch_object_detection/yolov3_spp/build_utils/utils.py
@@ -378,7 +378,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6,
         # Batched NMS
         c = x[:, 5] * 0 if agnostic else x[:, 5]  # classes
         boxes, scores = x[:, :4].clone() + c.view(-1, 1) * max_wh, x[:, 4]  # boxes (offset by class), scores
-        i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
+        i = torchvision.ops.nms(boxes, scores, iou_thres)
         i = i[:max_num]  # 最多只保留前max_num个目标信息
         if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
             try:  # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)