updates
This commit is contained in:
parent
cf2caaad41
commit
e5d11c68ac
12
detect.py
12
detect.py
|
@ -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))
|
||||||
|
|
|
@ -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 = []
|
||||||
|
|
||||||
|
|
2
test.py
2
test.py
|
@ -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):
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue