updates
This commit is contained in:
parent
0bac735cc6
commit
cf2caaad41
|
@ -67,7 +67,7 @@ def detect(
|
||||||
|
|
||||||
if detections is not None and len(detections) > 0:
|
if detections is not None and len(detections) > 0:
|
||||||
# Rescale boxes from 416 to true image size
|
# Rescale boxes from 416 to true image size
|
||||||
scale_coords(img_size, detections[:, :4], im0.shape).round()
|
detections[:, :4] = scale_coords(img_size, detections[:, :4], im0.shape).round()
|
||||||
|
|
||||||
# Print results to screen
|
# Print results to screen
|
||||||
for c in detections[:, -1].unique():
|
for c in detections[:, -1].unique():
|
||||||
|
|
14
models.py
14
models.py
|
@ -121,12 +121,12 @@ class YOLOLayer(nn.Module):
|
||||||
if ONNX_EXPORT:
|
if ONNX_EXPORT:
|
||||||
bs = 1 # batch size
|
bs = 1 # batch size
|
||||||
else:
|
else:
|
||||||
bs, nx, ny = p.shape[0], p.shape[-2], p.shape[-1]
|
bs, ny, nx = p.shape[0], p.shape[-2], p.shape[-1]
|
||||||
if (self.nx, self.ny) != (nx, ny):
|
if (self.ny, self.nx) != (ny, nx):
|
||||||
create_grids(self, img_size, (nx, ny), p.device)
|
create_grids(self, img_size, (ny, nx), p.device)
|
||||||
|
|
||||||
# p.view(bs, 255, 13, 13) -- > (bs, 3, 13, 13, 85) # (bs, anchors, grid, grid, classes + xywh)
|
# p.view(bs, 255, 13, 13) -- > (bs, 3, 13, 13, 85) # (bs, anchors, grid, grid, classes + xywh)
|
||||||
p = p.view(bs, self.na, self.nc + 5, self.nx, self.ny).permute(0, 1, 3, 4, 2).contiguous() # prediction
|
p = p.view(bs, self.na, self.nc + 5, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous() # prediction
|
||||||
|
|
||||||
if self.training:
|
if self.training:
|
||||||
return p
|
return p
|
||||||
|
@ -235,13 +235,13 @@ def get_yolo_layers(model):
|
||||||
|
|
||||||
|
|
||||||
def create_grids(self, img_size, ng, device='cpu'):
|
def create_grids(self, img_size, ng, device='cpu'):
|
||||||
nx, ny = ng # x and y grid size
|
ny, nx = ng # x and y grid size
|
||||||
self.img_size = img_size
|
self.img_size = img_size
|
||||||
self.stride = img_size / max(ng)
|
self.stride = img_size / max(ng)
|
||||||
|
|
||||||
# build xy offsets
|
# build xy offsets
|
||||||
yv, xv = torch.meshgrid([torch.arange(nx), torch.arange(ny)])
|
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
|
||||||
self.grid_xy = torch.stack((xv, yv), 2).to(device).float().view((1, 1, nx, ny, 2))
|
self.grid_xy = torch.stack((xv, yv), 2).to(device).float().view((1, 1, ny, nx, 2))
|
||||||
|
|
||||||
# build wh gains
|
# build wh gains
|
||||||
self.anchor_vec = self.anchors.to(device) / self.stride
|
self.anchor_vec = self.anchors.to(device) / self.stride
|
||||||
|
|
|
@ -74,7 +74,7 @@ class LoadImages: # for inference
|
||||||
print('image %g/%g %s: ' % (self.count, self.nF, path), end='')
|
print('image %g/%g %s: ' % (self.count, self.nF, path), end='')
|
||||||
|
|
||||||
# Padded resize
|
# Padded resize
|
||||||
img, _, _, _ = letterbox_rect(img0, height=self.height)
|
img, _, _, _ = letterbox(img0, height=self.height)
|
||||||
print('%gx%g ' % img.shape[:2], end='') # print image size
|
print('%gx%g ' % img.shape[:2], end='') # print image size
|
||||||
|
|
||||||
# Normalize RGB
|
# Normalize RGB
|
||||||
|
|
|
@ -103,8 +103,10 @@ def xywh2xyxy(x):
|
||||||
def scale_coords(img1_shape, coords, img0_shape):
|
def scale_coords(img1_shape, coords, img0_shape):
|
||||||
# Rescale coords1 (xyxy) from img1_shape to img0_shape
|
# Rescale coords1 (xyxy) from img1_shape to img0_shape
|
||||||
gain = img1_shape / max(img0_shape[:2]) # gain = old / new
|
gain = img1_shape / max(img0_shape[:2]) # gain = old / new
|
||||||
pad_x = np.mod(img1_shape - img0_shape[1] * gain, 32) / 2 # width padding
|
# pad_x = np.mod(img1_shape - img0_shape[1] * gain, 32) / 2 # width padding
|
||||||
pad_y = np.mod(img1_shape - img0_shape[0] * gain, 32) / 2 # height padding
|
# pad_y = np.mod(img1_shape - img0_shape[0] * gain, 32) / 2 # height padding
|
||||||
|
pad_x = (img1_shape - img0_shape[1] * gain) / 2 # width padding
|
||||||
|
pad_y = (img1_shape - img0_shape[0] * gain) / 2 # height padding
|
||||||
coords[:, [0, 2]] -= pad_x
|
coords[:, [0, 2]] -= pad_x
|
||||||
coords[:, [1, 3]] -= pad_y
|
coords[:, [1, 3]] -= pad_y
|
||||||
coords[:, :4] /= gain
|
coords[:, :4] /= gain
|
||||||
|
|
Loading…
Reference in New Issue