Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-02-10 21:41:57 +01:00
parent 715c4575bf
commit 62761cffe6
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
detect.py
View File

@ -38,8 +38,8 @@ def detect(
# Set Dataloader # Set Dataloader
dataloader = load_images(images, img_size=img_size) dataloader = load_images(images, img_size=img_size)
# Classes and colors # Get classes and colors
classes = load_classes(parse_data_cfg('cfg/coco.data')['names']) # Extracts class labels from file classes = load_classes(parse_data_cfg('cfg/coco.data')['names'])
colors = [[random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)] for _ in range(len(classes))] colors = [[random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)] for _ in range(len(classes))]
for i, (path, img, im0) in enumerate(dataloader): for i, (path, img, im0) in enumerate(dataloader):
@ -49,7 +49,7 @@ def detect(
# Get detections # Get detections
img = torch.from_numpy(img).unsqueeze(0).to(device) img = torch.from_numpy(img).unsqueeze(0).to(device)
if ONNX_EXPORT: if ONNX_EXPORT:
pred = torch.onnx._export(model, img, 'weights/model.onnx', verbose=True) torch.onnx._export(model, img, 'weights/model.onnx', verbose=True)
return # ONNX export return # ONNX export
pred = model(img) pred = model(img)
pred = pred[pred[:, :, 4] > conf_thres] pred = pred[pred[:, :, 4] > conf_thres]