This commit is contained in:
Glenn Jocher 2019-04-22 14:31:23 +02:00
parent 0bac735cc6
commit cf2caaad41
4 changed files with 13 additions and 11 deletions

View File

@ -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():

View File

@ -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

View File

@ -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

View File

@ -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