Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-01-06 14:23:04 +02:00
parent 178e1a346b
commit 6e1ff541c9
2 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
models.py
View File

@ -254,9 +254,14 @@ class YOLOLayer(nn.Module):
xy = torch.sigmoid(p[..., 0:2]) + self.grid_xy # x, y
width_height = torch.exp(p[..., 2:4]) * self.anchor_wh # width, height
p_conf = torch.sigmoid(p[..., 4:5]) # Conf
## p_cls = torch.sigmoid(p[..., 5:85]) # Class
p_cls = F.softmax(p[..., 5:85], 2) * p_conf # SSD-like conf
# p_cls = torch.exp(p[..., 5:85]) / torch.exp(p[..., 5:85]).sum(2).unsqueeze(2) #* p_conf # F.softmax() equivalent
p_cls = p[..., 5:85]
# Broadcasting only supported on first dimension in CoreML. See onnx-coreml/_operators.py
# p_cls = F.softmax(p_cls, 2) * p_conf # SSD-like conf
p_cls = torch.exp(p_cls).permute(2, 1, 0)
p_cls = p_cls / p_cls.sum(0).unsqueeze(0) * p_conf.permute(2, 1, 0) # F.softmax() equivalent
p_cls = p_cls.permute(2, 1, 0)
return torch.cat((xy / nG, width_height, p_conf, p_cls), 2).squeeze().t()
p[..., 0] = torch.sigmoid(p[..., 0]) + self.grid_x # x

11
utils/onnx2coreml.py
View File

@ -68,13 +68,14 @@ def main():
spec.neuralNetwork.preprocessing[0].featureName = '0'
yolov3_model.save(name + '.mlmodel')
# yolov3_model.visualize_spec()
print(spec.description)
# 2.5. Try to Predict:
from PIL import Image
img = Image.open('../yolov3/data/samples/zidane_416.jpg')
out = yolov3_model.predict({'0': img})
print(out['141'].shape, out['143'].shape)
out = yolov3_model.predict({'0': img}, useCPUOnly=True)
print(out['148'].shape, out['150'].shape)
# 3. Create NMS protobuf
import numpy as np
@ -106,15 +107,15 @@ def main():
del ma_type.shape[:]
nms = nms_spec.nonMaximumSuppression
nms.confidenceInputFeatureName = '141' # 1x507x80
nms.coordinatesInputFeatureName = '143' # 1x507x4
nms.confidenceInputFeatureName = '148' # 1x507x80
nms.coordinatesInputFeatureName = '150' # 1x507x4
nms.confidenceOutputFeatureName = 'confidence'
nms.coordinatesOutputFeatureName = 'coordinates'
nms.iouThresholdInputFeatureName = 'iouThreshold'
nms.confidenceThresholdInputFeatureName = 'confidenceThreshold'
nms.iouThreshold = 0.6
nms.confidenceThreshold = 0.9
nms.confidenceThreshold = 0.3
nms.pickTop.perClass = True
labels = np.loadtxt('../yolov3/data/coco.names', dtype=str, delimiter='\n')