Forráskód Böngészése

train YOLOv2 on VOC

yjh0410 2 éve
szülő
commit
15c4626136
3 módosított fájl, 14 hozzáadás és 4 törlés
  1. 6 1
      models/yolov1/matcher.py
  2. 7 2
      models/yolov2/matcher.py
  3. 1 1
      train.sh

+ 6 - 1
models/yolov1/matcher.py

@@ -47,8 +47,13 @@ class YoloMatcher(object):
                 grid_y = int(ys_c)
 
                 if grid_x < fmp_w and grid_y < fmp_h:
+                    # obj
                     gt_objectness[batch_index, grid_y, grid_x] = 1.0
-                    gt_classes[batch_index, grid_y, grid_x, int(gt_label)] = 1.0
+                    # cls
+                    cls_ont_hot = np.zeros(self.num_classes)
+                    cls_ont_hot[int(gt_label)] = 1.0
+                    gt_classes[batch_index, grid_y, grid_x] = cls_ont_hot
+                    # box
                     gt_bboxes[batch_index, grid_y, grid_x] = np.array([x1, y1, x2, y2])
 
         # [B, M, C]

+ 7 - 2
models/yolov2/matcher.py

@@ -83,7 +83,7 @@ class Yolov2Matcher(object):
 
                 # check
                 if bw < 1. or bh < 1.:
-                    return False    
+                    continue    
 
                 # compute IoU
                 iou = self.compute_iou(self.anchor_boxes, gt_box)
@@ -117,8 +117,13 @@ class Yolov2Matcher(object):
                 for result in label_assignment_results:
                     grid_x, grid_y, anchor_idx = result
                     if grid_x < fmp_w and grid_y < fmp_h:
+                        # obj
                         gt_objectness[batch_index, grid_y, grid_x, anchor_idx] = 1.0
-                        gt_classes[batch_index, grid_y, grid_x, anchor_idx, int(gt_label)] = 1.0
+                        # cls
+                        cls_ont_hot = np.zeros(self.num_classes)
+                        cls_ont_hot[int(gt_label)] = 1.0
+                        gt_classes[batch_index, grid_y, grid_x, anchor_idx] = cls_ont_hot
+                        # box
                         gt_bboxes[batch_index, grid_y, grid_x, anchor_idx] = np.array([x1, y1, x2, y2])
 
         # [B, H, W, A, C] -> [B, HWA, C]

+ 1 - 1
train.sh

@@ -3,7 +3,7 @@ python train.py \
         --cuda \
         -d voc \
         --root /mnt/share/ssd2/dataset/ \
-        -m yolov1 \
+        -m yolov2 \
         -bs 16 \
         -size 640 \
         --wp_epoch 1 \