|
import argparse |
|
|
|
from utils.datasets import * |
|
from utils.utils import * |
|
|
|
ONNX_EXPORT = False |
|
|
|
|
|
def detect(save_img=False): |
|
out, source, weights, half, view_img, save_txt, imgsz = \ |
|
opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt, opt.img_size |
|
webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt') |
|
|
|
|
|
device = torch_utils.select_device(opt.device) |
|
if os.path.exists(out): |
|
shutil.rmtree(out) |
|
os.makedirs(out) |
|
|
|
|
|
google_utils.attempt_download(weights) |
|
model = torch.load(weights, map_location=device)['model'] |
|
|
|
|
|
|
|
classify = False |
|
if classify: |
|
modelc = torch_utils.load_classifier(name='resnet101', n=2) |
|
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) |
|
modelc.to(device).eval() |
|
|
|
|
|
model.to(device).eval() |
|
|
|
|
|
|
|
|
|
|
|
half = half and device.type != 'cpu' |
|
if half: |
|
model.half() |
|
|
|
|
|
vid_path, vid_writer = None, None |
|
if webcam: |
|
view_img = True |
|
torch.backends.cudnn.benchmark = True |
|
dataset = LoadStreams(source, img_size=imgsz) |
|
else: |
|
save_img = True |
|
dataset = LoadImages(source, img_size=imgsz) |
|
|
|
|
|
names = model.names |
|
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))] |
|
|
|
|
|
t0 = time.time() |
|
img = torch.zeros((1, 3, imgsz, imgsz), device=device) |
|
_ = model(img.half() if half else img.float()) if device.type != 'cpu' else None |
|
for path, img, im0s, vid_cap in dataset: |
|
img = torch.from_numpy(img).to(device) |
|
img = img.half() if half else img.float() |
|
img /= 255.0 |
|
if img.ndimension() == 3: |
|
img = img.unsqueeze(0) |
|
|
|
|
|
t1 = torch_utils.time_synchronized() |
|
pred = model(img, augment=opt.augment)[0] |
|
t2 = torch_utils.time_synchronized() |
|
|
|
|
|
if half: |
|
pred = pred.float() |
|
|
|
|
|
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, |
|
multi_label=False, classes=opt.classes, agnostic=opt.agnostic_nms) |
|
|
|
|
|
if classify: |
|
pred = apply_classifier(pred, modelc, img, im0s) |
|
|
|
|
|
for i, det in enumerate(pred): |
|
if webcam: |
|
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy() |
|
else: |
|
p, s, im0 = path, '', im0s |
|
|
|
save_path = str(Path(out) / Path(p).name) |
|
s += '%gx%g ' % img.shape[2:] |
|
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] |
|
if det is not None and len(det): |
|
|
|
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() |
|
|
|
|
|
for c in det[:, -1].unique(): |
|
n = (det[:, -1] == c).sum() |
|
s += '%g %ss, ' % (n, names[int(c)]) |
|
|
|
|
|
for *xyxy, conf, cls in det: |
|
if save_txt: |
|
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() |
|
with open(save_path[:save_path.rfind('.')] + '.txt', 'a') as file: |
|
file.write(('%g ' * 5 + '\n') % (cls, *xywh)) |
|
|
|
if save_img or view_img: |
|
label = '%s %.2f' % (names[int(cls)], conf) |
|
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) |
|
|
|
|
|
print('%sDone. (%.3fs)' % (s, t2 - t1)) |
|
|
|
|
|
if view_img: |
|
cv2.imshow(p, im0) |
|
if cv2.waitKey(1) == ord('q'): |
|
raise StopIteration |
|
|
|
|
|
if save_img: |
|
if dataset.mode == 'images': |
|
cv2.imwrite(save_path, im0) |
|
else: |
|
if vid_path != save_path: |
|
vid_path = save_path |
|
if isinstance(vid_writer, cv2.VideoWriter): |
|
vid_writer.release() |
|
|
|
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) |
|
|
|
if save_txt or save_img: |
|
print('Results saved to %s' % os.getcwd() + os.sep + out) |
|
if platform == 'darwin': |
|
os.system('open ' + save_path) |
|
|
|
print('Done. (%.3fs)' % (time.time() - t0)) |
|
|
|
|
|
if __name__ == '__main__': |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='model.pt path') |
|
parser.add_argument('--source', type=str, default='inference/images', help='source') |
|
parser.add_argument('--output', type=str, default='inference/output', help='output folder') |
|
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') |
|
parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold') |
|
parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') |
|
parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)') |
|
parser.add_argument('--half', action='store_true', help='half precision FP16 inference') |
|
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') |
|
parser.add_argument('--view-img', action='store_true', help='display results') |
|
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') |
|
parser.add_argument('--classes', nargs='+', type=int, help='filter by class') |
|
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') |
|
parser.add_argument('--augment', action='store_true', help='augmented inference') |
|
opt = parser.parse_args() |
|
print(opt) |
|
|
|
with torch.no_grad(): |
|
detect() |
|
|