car-detection-bayes/detect.py

179 lines
7.5 KiB
Python
Raw Normal View History

2018-08-26 08:51:39 +00:00
import argparse
2019-02-12 15:58:07 +00:00
from sys import platform
2018-08-26 08:51:39 +00:00
from models import * # set ONNX_EXPORT in models.py
2018-08-26 08:51:39 +00:00
from utils.datasets import *
from utils.utils import *
2019-01-08 18:37:23 +00:00
2020-01-02 17:50:11 +00:00
def detect(save_img=False):
2019-08-31 17:11:59 +00:00
img_size = (320, 192) if ONNX_EXPORT else opt.img_size # (320, 192) or (416, 256) or (608, 352) for (height, width)
2020-01-02 17:50:11 +00:00
out, source, weights, half, view_img, save_txt = opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt
2019-09-19 16:43:29 +00:00
webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
2019-08-31 16:58:30 +00:00
2019-07-15 15:00:04 +00:00
# Initialize
2019-09-13 14:00:52 +00:00
device = torch_utils.select_device(device='cpu' if ONNX_EXPORT else opt.device)
2019-08-31 16:58:30 +00:00
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
2018-08-26 08:51:39 +00:00
2019-02-11 11:32:54 +00:00
# Initialize model
2019-08-31 17:11:59 +00:00
model = Darknet(opt.cfg, img_size)
2018-08-26 08:51:39 +00:00
2019-02-11 11:32:54 +00:00
# Load weights
2019-09-19 16:09:16 +00:00
attempt_download(weights)
if weights.endswith('.pt'): # pytorch format
model.load_state_dict(torch.load(weights, map_location=device)['model'])
2018-12-06 12:01:49 +00:00
else: # darknet format
2020-01-17 18:55:30 +00:00
load_darknet_weights(model, weights)
2018-08-26 08:51:39 +00:00
2019-10-10 20:54:20 +00:00
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101', n=2) # initialize
2019-10-12 11:59:07 +00:00
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights
2019-10-11 23:18:41 +00:00
modelc.to(device).eval()
2019-10-10 20:54:20 +00:00
2019-04-20 20:46:23 +00:00
# Fuse Conv2d + BatchNorm2d layers
2019-07-30 10:39:17 +00:00
# model.fuse()
2020-01-30 20:12:04 +00:00
# torch_utils.model_info(model, report='summary') # 'full' or 'summary'
2019-04-19 18:41:18 +00:00
2019-04-22 14:21:21 +00:00
# Eval mode
2018-08-26 08:51:39 +00:00
model.to(device).eval()
2019-07-30 10:39:17 +00:00
# Export mode
2019-04-22 14:21:21 +00:00
if ONNX_EXPORT:
2020-01-30 20:12:04 +00:00
model.fuse()
2019-08-31 17:11:59 +00:00
img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192)
2020-02-15 05:32:29 +00:00
torch.onnx.export(model, img, 'weights/export.onnx', verbose=False, opset_version=11)
2019-11-08 01:55:00 +00:00
# Validate exported model
import onnx
model = onnx.load('weights/export.onnx') # Load the ONNX model
onnx.checker.check_model(model) # Check that the IR is well formed
print(onnx.helper.printable_graph(model.graph)) # Print a human readable representation of the graph
2019-04-22 14:21:21 +00:00
return
2019-07-31 22:08:28 +00:00
# Half precision
half = half and device.type != 'cpu' # half precision only supported on CUDA
if half:
2019-07-31 22:08:28 +00:00
model.half()
2018-08-26 08:51:39 +00:00
# Set Dataloader
2019-04-02 11:43:18 +00:00
vid_path, vid_writer = None, None
2019-09-19 16:43:29 +00:00
if webcam:
2019-09-20 13:24:00 +00:00
view_img = True
2019-09-10 12:25:56 +00:00
torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference
2020-02-23 02:18:38 +00:00
dataset = LoadStreams(source, img_size=img_size)
2019-02-11 12:45:04 +00:00
else:
save_img = True
2020-02-23 02:18:38 +00:00
dataset = LoadImages(source, img_size=img_size)
2019-02-08 21:43:05 +00:00
2019-12-15 20:47:53 +00:00
# Get names and colors
names = load_classes(opt.names)
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
2018-08-26 08:51:39 +00:00
2019-07-31 22:08:28 +00:00
# Run inference
2019-08-01 00:21:40 +00:00
t0 = time.time()
2019-09-09 23:34:23 +00:00
for path, img, im0s, vid_cap in dataset:
2019-02-08 21:43:05 +00:00
t = time.time()
2019-09-09 23:34:23 +00:00
img = torch.from_numpy(img).to(device)
2020-02-23 02:18:38 +00:00
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
2019-09-09 23:34:23 +00:00
if img.ndimension() == 3:
img = img.unsqueeze(0)
2020-02-23 02:18:38 +00:00
# Inference
pred = model(img)[0].float() if half else model(img)[0]
2019-09-27 21:46:45 +00:00
2019-10-10 20:54:20 +00:00
# Apply NMS
2020-01-10 17:30:05 +00:00
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
2019-10-10 20:54:20 +00:00
2019-12-23 19:25:15 +00:00
# Apply Classifier
2019-10-10 20:54:20 +00:00
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
2019-09-19 16:43:29 +00:00
if webcam: # batch_size >= 1
2019-09-09 23:34:23 +00:00
p, s, im0 = path[i], '%g: ' % i, im0s[i]
2019-04-28 21:16:21 +00:00
else:
2019-09-09 23:34:23 +00:00
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
s += '%gx%g ' % img.shape[2:] # print string
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
2019-12-15 20:47:53 +00:00
s += '%g %ss, ' % (n, names[int(c)]) # add to string
2019-09-09 23:34:23 +00:00
# Write results
2019-12-20 02:09:13 +00:00
for *xyxy, conf, cls in det:
2019-09-09 23:34:23 +00:00
if save_txt: # Write to file
with open(save_path + '.txt', 'a') as file:
file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf))
2019-09-19 13:31:28 +00:00
if save_img or view_img: # Add bbox to image
2019-12-15 20:47:53 +00:00
label = '%s %.2f' % (names[int(cls)], conf)
2019-09-09 23:34:23 +00:00
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)])
2020-01-01 20:44:33 +00:00
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, time.time() - t))
2019-09-09 23:34:23 +00:00
# Stream results
2019-09-19 13:31:28 +00:00
if view_img:
2019-09-09 23:34:23 +00:00
cv2.imshow(p, im0)
2019-12-12 21:56:56 +00:00
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
2019-09-09 23:34:23 +00:00
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h))
vid_writer.write(im0)
2019-04-02 11:43:18 +00:00
if save_txt or save_img:
2019-08-31 16:58:30 +00:00
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + out + ' ' + save_path)
2018-11-21 18:24:00 +00:00
2019-08-01 00:21:40 +00:00
print('Done. (%.3fs)' % (time.time() - t0))
2018-08-26 08:51:39 +00:00
if __name__ == '__main__':
parser = argparse.ArgumentParser()
2019-12-15 20:47:53 +00:00
parser.add_argument('--cfg', type=str, default='cfg/yolov3-spp.cfg', help='*.cfg path')
parser.add_argument('--names', type=str, default='data/coco.names', help='*.names path')
2020-02-17 07:12:07 +00:00
parser.add_argument('--weights', type=str, default='weights/yolov3-spp-ultralytics.pt', help='weights path')
2019-09-20 13:23:08 +00:00
parser.add_argument('--source', type=str, default='data/samples', help='source') # input file/folder, 0 for webcam
2019-08-31 16:58:30 +00:00
parser.add_argument('--output', type=str, default='output', help='output folder') # output folder
2019-04-29 15:49:09 +00:00
parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
2019-08-05 11:57:18 +00:00
parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
2020-02-17 07:30:14 +00:00
parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
2019-08-31 16:58:30 +00:00
parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
2019-07-31 22:08:28 +00:00
parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
2019-09-13 14:00:52 +00:00
parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu')
2019-09-19 16:54:16 +00:00
parser.add_argument('--view-img', action='store_true', help='display results')
2020-01-15 18:22:59 +00:00
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
2020-01-01 20:44:33 +00:00
parser.add_argument('--classes', nargs='+', type=int, help='filter by class')
2020-01-10 17:30:05 +00:00
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
opt = parser.parse_args()
print(opt)
2019-02-10 20:06:22 +00:00
with torch.no_grad():
2019-08-31 16:58:30 +00:00
detect()