This commit is contained in:
Glenn Jocher 2019-04-22 14:59:39 +02:00
parent cf2caaad41
commit e5d11c68ac
5 changed files with 12 additions and 16 deletions

View File

@ -63,19 +63,19 @@ def detect(
# Get detections # Get detections
img = torch.from_numpy(img).unsqueeze(0).to(device) img = torch.from_numpy(img).unsqueeze(0).to(device)
pred, _ = model(img) pred, _ = model(img)
detections = non_max_suppression(pred, conf_thres, nms_thres)[0] det = non_max_suppression(pred, conf_thres, nms_thres)[0]
if detections is not None and len(detections) > 0: if det is not None and len(det) > 0:
# Rescale boxes from 416 to true image size # Rescale boxes from 416 to true image size
detections[:, :4] = scale_coords(img_size, detections[:, :4], im0.shape).round() det[:, :4] = scale_coords(img.shape, det[:, :4], im0.shape).round()
# Print results to screen # Print results to screen
for c in detections[:, -1].unique(): for c in det[:, -1].unique():
n = (detections[:, -1] == c).sum() n = (det[:, -1] == c).sum()
print('%g %ss' % (n, classes[int(c)]), end=', ') print('%g %ss' % (n, classes[int(c)]), end=', ')
# Draw bounding boxes and labels of detections # Draw bounding boxes and labels of detections
for *xyxy, conf, cls_conf, cls in detections: for *xyxy, conf, cls_conf, cls in det:
if save_txt: # Write to file if save_txt: # Write to file
with open(save_path + '.txt', 'a') as file: with open(save_path + '.txt', 'a') as file:
file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf)) file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf))

View File

@ -182,7 +182,7 @@ class Darknet(nn.Module):
self.yolo_layers = get_yolo_layers(self) self.yolo_layers = get_yolo_layers(self)
def forward(self, x, var=None): def forward(self, x, var=None):
img_size = x.shape[-1] img_size = max(x.shape[-2:])
layer_outputs = [] layer_outputs = []
output = [] output = []

View File

@ -93,7 +93,7 @@ def test(
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = int(Path(paths[si]).stem.split('_')[-1]) image_id = int(Path(paths[si]).stem.split('_')[-1])
box = pred[:, :4].clone() # xyxy box = pred[:, :4].clone() # xyxy
scale_coords(img_size, box, shapes[si]) # to original shape scale_coords(imgs[si].shape, box, shapes[si]) # to original shape
box = xyxy2xywh(box) # xywh box = xyxy2xywh(box) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for di, d in enumerate(pred): for di, d in enumerate(pred):

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(img0, height=self.height) img, _, _, _ = letterbox_rect(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

@ -102,13 +102,9 @@ 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 = max(img1_shape) / max(img0_shape) # gain = old / new
# pad_x = np.mod(img1_shape - img0_shape[1] * gain, 32) / 2 # width padding coords[:, [0, 2]] -= (img1_shape[3] - img0_shape[1] * gain) / 2 # x padding
# pad_y = np.mod(img1_shape - img0_shape[0] * gain, 32) / 2 # height padding coords[:, [1, 3]] -= (img1_shape[2] - img0_shape[0] * gain) / 2 # y 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[:, [1, 3]] -= pad_y
coords[:, :4] /= gain coords[:, :4] /= gain
coords[:, :4] = coords[:, :4].clamp(min=0) coords[:, :4] = coords[:, :4].clamp(min=0)
return coords return coords