car-detection-bayes/detect.py

127 lines
4.5 KiB
Python
Raw Normal View History

2018-08-26 08:51:39 +00:00
import argparse
2019-02-12 17:05:58 +00:00
import shutil
2018-08-26 08:51:39 +00:00
import time
2019-02-12 17:05:58 +00:00
from pathlib import Path
2019-02-12 15:58:07 +00:00
from sys import platform
2018-08-26 08:51:39 +00:00
from models import *
from utils.datasets import *
from utils.utils import *
2019-01-08 18:37:23 +00:00
2019-02-10 20:06:22 +00:00
def detect(
cfg,
weights,
images,
2019-02-17 16:30:16 +00:00
output='output', # output folder
2019-02-10 20:06:22 +00:00
img_size=416,
conf_thres=0.3,
nms_thres=0.45,
save_txt=False,
2019-02-11 12:45:04 +00:00
save_images=True,
2019-02-11 13:19:35 +00:00
webcam=False
2019-02-10 20:06:22 +00:00
):
device = torch_utils.select_device()
2019-02-17 16:30:16 +00:00
if os.path.exists(output):
shutil.rmtree(output) # delete output folder
2019-02-12 16:29:13 +00:00
os.makedirs(output) # make new output folder
2018-08-26 08:51:39 +00:00
2019-02-11 11:32:54 +00:00
# Initialize model
2019-02-08 21:43:05 +00:00
model = Darknet(cfg, img_size)
2018-08-26 08:51:39 +00:00
2019-02-11 11:32:54 +00:00
# Load weights
2019-02-08 21:43:05 +00:00
if weights.endswith('.pt'): # pytorch format
2019-02-17 16:34:45 +00:00
if weights.endswith('yolov3.pt') and not os.path.exists(weights):
if (platform == 'darwin') or (platform == 'linux'):
os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights)
2019-02-08 22:20:41 +00:00
model.load_state_dict(torch.load(weights, map_location='cpu')['model'])
2018-12-06 12:01:49 +00:00
else: # darknet format
2019-02-08 21:43:05 +00:00
load_darknet_weights(model, weights)
2018-08-26 08:51:39 +00:00
model.to(device).eval()
# Set Dataloader
2019-02-11 12:45:04 +00:00
if webcam:
save_images = False
2019-02-11 16:25:32 +00:00
dataloader = LoadWebcam(img_size=img_size)
2019-02-11 12:45:04 +00:00
else:
dataloader = LoadImages(images, img_size=img_size)
2019-02-08 21:43:05 +00:00
2019-02-10 20:41:57 +00:00
# Get classes and colors
classes = load_classes(parse_data_cfg('cfg/coco.data')['names'])
2019-02-08 22:08:26 +00:00
colors = [[random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)] for _ in range(len(classes))]
2018-08-26 08:51:39 +00:00
2019-02-09 18:24:51 +00:00
for i, (path, img, im0) in enumerate(dataloader):
2019-02-08 21:43:05 +00:00
t = time.time()
2019-02-11 17:15:51 +00:00
if webcam:
print('webcam frame %g: ' % (i + 1), end='')
else:
print('image %g/%g %s: ' % (i + 1, len(dataloader), path), end='')
2019-02-12 16:56:24 +00:00
save_path = str(Path(output) / Path(path).name)
2018-08-26 08:51:39 +00:00
# Get detections
2019-02-10 20:06:22 +00:00
img = torch.from_numpy(img).unsqueeze(0).to(device)
if ONNX_EXPORT:
2019-02-11 17:15:51 +00:00
torch.onnx.export(model, img, 'weights/model.onnx', verbose=True)
2019-02-11 17:17:38 +00:00
return
2019-02-10 20:06:22 +00:00
pred = model(img)
2019-02-11 12:45:04 +00:00
pred = pred[pred[:, :, 4] > conf_thres] # remove boxes < threshold
2018-08-26 08:51:39 +00:00
2019-02-10 20:06:22 +00:00
if len(pred) > 0:
2019-02-11 12:45:04 +00:00
# Run NMS on predictions
2019-02-10 20:06:22 +00:00
detections = non_max_suppression(pred.unsqueeze(0), conf_thres, nms_thres)[0]
2019-02-08 21:43:05 +00:00
2019-02-10 20:06:22 +00:00
# Rescale boxes from 416 to true image size
detections[:, :4] = scale_coords(img_size, detections[:, :4], im0.shape)
2019-02-08 21:43:05 +00:00
2019-02-11 12:45:04 +00:00
# Print results to screen
2019-02-10 20:06:22 +00:00
unique_classes = detections[:, -1].cpu().unique()
2019-02-11 17:15:51 +00:00
for c in unique_classes:
n = (detections[:, -1].cpu() == c).sum()
print('%g %ss' % (n, classes[int(c)]), end=', ')
2019-02-08 21:43:05 +00:00
2019-02-11 12:45:04 +00:00
# Draw bounding boxes and labels of detections
2019-02-11 11:26:30 +00:00
for x1, y1, x2, y2, conf, cls_conf, cls in detections:
2019-02-10 20:06:22 +00:00
if save_txt: # Write to file
2019-02-11 11:26:30 +00:00
with open(save_path + '.txt', 'a') as file:
2019-02-11 12:45:04 +00:00
file.write('%g %g %g %g %g %g\n' %
(x1, y1, x2, y2, cls, cls_conf * conf))
2019-02-08 21:43:05 +00:00
2019-02-11 12:45:04 +00:00
# Add bbox to the image
label = '%s %.2f' % (classes[int(cls)], conf)
plot_one_box([x1, y1, x2, y2], im0, label=label, color=colors[int(cls)])
2018-08-26 08:51:39 +00:00
2019-02-11 13:19:06 +00:00
dt = time.time() - t
print('Done. (%.3fs)' % dt)
2018-08-26 08:51:39 +00:00
2019-02-11 12:45:04 +00:00
if save_images: # Save generated image with detections
cv2.imwrite(save_path, im0)
if webcam: # Show live webcam
2019-02-11 13:19:06 +00:00
cv2.imshow(weights + ' - %.2f FPS' % (1 / dt), im0)
2019-02-11 12:45:04 +00:00
2019-02-12 17:21:06 +00:00
if save_images and (platform == 'darwin'): # linux/macos
2019-02-12 17:07:23 +00:00
os.system('open ' + output + ' ' + save_path)
2018-11-21 18:24:00 +00:00
2018-08-26 08:51:39 +00:00
if __name__ == '__main__':
parser = argparse.ArgumentParser()
2019-02-11 12:47:58 +00:00
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
2019-02-19 14:01:08 +00:00
parser.add_argument('--weights', type=str, default='weights/yolov3.weights', help='path to weights file')
2019-02-08 22:28:00 +00:00
parser.add_argument('--images', type=str, default='data/samples', help='path to images')
parser.add_argument('--img-size', type=int, default=32 * 13, help='size of each image dimension')
parser.add_argument('--conf-thres', type=float, default=0.50, help='object confidence threshold')
parser.add_argument('--nms-thres', type=float, default=0.45, help='iou threshold for non-maximum suppression')
opt = parser.parse_args()
print(opt)
2019-02-10 20:06:22 +00:00
with torch.no_grad():
detect(
opt.cfg,
opt.weights,
opt.images,
img_size=opt.img_size,
conf_thres=opt.conf_thres,
nms_thres=opt.nms_thres
)