Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Yolo V2 support #561

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 12 additions & 17 deletions tool/darknet2pytorch.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tool.region_loss import RegionLoss
from tool.yolo_layer import YoloLayer
from tool.config import *
from tool.torch_utils import *
Expand Down Expand Up @@ -208,12 +207,8 @@ def forward(self, x):
x = x1 * x2
outputs[ind] = x
elif block['type'] == 'region':
continue
if self.loss:
self.loss = self.loss + self.models[ind](x)
else:
self.loss = self.models[ind](x)
outputs[ind] = None
boxes = self.models[ind](x)
out_boxes.append(boxes)
elif block['type'] == 'yolo':
# if self.training:
# pass
Expand Down Expand Up @@ -392,19 +387,19 @@ def create_network(self, blocks):
out_strides.append(prev_stride)
models.append(model)
elif block['type'] == 'region':
loss = RegionLoss()
region = YoloLayer()
anchors = block['anchors'].split(',')
loss.anchors = [float(i) for i in anchors]
loss.num_classes = int(block['classes'])
loss.num_anchors = int(block['num'])
loss.anchor_step = len(loss.anchors) // loss.num_anchors
loss.object_scale = float(block['object_scale'])
loss.noobject_scale = float(block['noobject_scale'])
loss.class_scale = float(block['class_scale'])
loss.coord_scale = float(block['coord_scale'])
region.anchors = [float(i) for i in anchors]
region.num_classes = int(block['classes'])
region.num_anchors = int(block['num'])
region.anchor_step = len(region.anchors) // region.num_anchors
region.scale_x_y = 1.0 # thre is not such value in region config
region.anchor_mask = [int(i) for i in range(len(anchors) // 2)] # region has no anchor masks
region.stride = 1 # not implemented for region
region.multiply_confs = False # do not multiply detection and class confidence
out_filters.append(prev_filters)
out_strides.append(prev_stride)
models.append(loss)
models.append(region)
elif block['type'] == 'yolo':
yolo_layer = YoloLayer()
anchors = block['anchors'].split(',')
Expand Down
4 changes: 2 additions & 2 deletions tool/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,8 +137,8 @@ def get_color(c, x, max_val):
t_size = cv2.getTextSize(msg, 0, 0.7, thickness=bbox_thick // 2)[0]
c1, c2 = (x1,y1), (x2, y2)
c3 = (c1[0] + t_size[0], c1[1] - t_size[1] - 3)
cv2.rectangle(img, (x1,y1), (np.float32(c3[0]), np.float32(c3[1])), rgb, -1)
img = cv2.putText(img, msg, (c1[0], np.float32(c1[1] - 2)), cv2.FONT_HERSHEY_SIMPLEX,0.7, (0,0,0), bbox_thick//2,lineType=cv2.LINE_AA)
cv2.rectangle(img, (x1,y1), (int(c3[0]), int(c3[1])), rgb, -1)
img = cv2.putText(img, msg, (c1[0], int(c1[1] - 2)), cv2.FONT_HERSHEY_SIMPLEX,0.7, (0,0,0), bbox_thick//2,lineType=cv2.LINE_AA)

img = cv2.rectangle(img, (x1, y1), (x2, y2), rgb, bbox_thick)
if savename:
Expand Down
13 changes: 8 additions & 5 deletions tool/yolo_layer.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ def yolo_forward(output, conf_thresh, num_classes, anchors, num_anchors, scale_x


def yolo_forward_dynamic(output, conf_thresh, num_classes, anchors, num_anchors, scale_x_y, only_objectness=1,
validation=False):
validation=False, multiply_confs=True):
# Output would be invalid if it does not satisfy this assert
# assert (output.size(1) == (5 + num_classes) * num_anchors)

Expand Down Expand Up @@ -280,7 +280,10 @@ def yolo_forward_dynamic(output, conf_thresh, num_classes, anchors, num_anchors,
# det_confs: [batch, num_anchors * H * W]

det_confs = det_confs.view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
confs = cls_confs * det_confs
if multiply_confs:
confs = cls_confs * det_confs
else:
confs = det_confs

# boxes: [batch, num_anchors * H * W, 1, 4]
# confs: [batch, num_anchors * H * W, num_classes]
Expand All @@ -292,7 +295,7 @@ class YoloLayer(nn.Module):
model_out: while inference,is post-processing inside or outside the model
true:outside
'''
def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1, stride=32, model_out=False):
def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1, stride=32, model_out=False, multiply_confs=True):
super(YoloLayer, self).__init__()
self.anchor_mask = anchor_mask
self.num_classes = num_classes
Expand All @@ -307,6 +310,7 @@ def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1, str
self.stride = stride
self.seen = 0
self.scale_x_y = 1
self.multiply_confs = multiply_confs

self.model_out = model_out

Expand All @@ -318,5 +322,4 @@ def forward(self, output, target=None):
masked_anchors += self.anchors[m * self.anchor_step:(m + 1) * self.anchor_step]
masked_anchors = [anchor / self.stride for anchor in masked_anchors]

return yolo_forward_dynamic(output, self.thresh, self.num_classes, masked_anchors, len(self.anchor_mask),scale_x_y=self.scale_x_y)

return yolo_forward_dynamic(output, self.thresh, self.num_classes, masked_anchors, len(self.anchor_mask),scale_x_y=self.scale_x_y, multiply_confs=self.multiply_confs)