Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Path: blob/master/detect.py
Views: 475
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license1"""2Run inference on images, videos, directories, streams, etc.34Usage:5$ python path/to/detect.py --weights yolov5s.pt --source 0 # webcam6img.jpg # image7vid.mp4 # video8path/ # directory9path/*.jpg # glob10'https://youtu.be/Zgi9g1ksQHc' # YouTube11'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream12"""1314import argparse15import os16import sys17from pathlib import Path1819import cv220import torch21import torch.backends.cudnn as cudnn2223FILE = Path(__file__).resolve()24ROOT = FILE.parents[0] # YOLOv5 root directory25if str(ROOT) not in sys.path:26sys.path.append(str(ROOT)) # add ROOT to PATH27ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative2829from models.common import DetectMultiBackend30from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams31from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,32increment_path, non_max_suppression, non_max_suppression_obb, print_args, scale_coords, scale_polys, strip_optimizer, xyxy2xywh)33from utils.plots import Annotator, colors, save_one_box34from utils.torch_utils import select_device, time_sync35from utils.rboxs_utils import poly2rbox, rbox2poly363738@torch.no_grad()39def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)40source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam41imgsz=(640, 640), # inference size (height, width)42conf_thres=0.25, # confidence threshold43iou_thres=0.45, # NMS IOU threshold44max_det=1000, # maximum detections per image45device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu46view_img=False, # show results47save_txt=False, # save results to *.txt48save_conf=False, # save confidences in --save-txt labels49save_crop=False, # save cropped prediction boxes50nosave=False, # do not save images/videos51classes=None, # filter by class: --class 0, or --class 0 2 352agnostic_nms=False, # class-agnostic NMS53augment=False, # augmented inference54visualize=False, # visualize features55update=False, # update all models56project=ROOT / 'runs/detect', # save results to project/name57name='exp', # save results to project/name58exist_ok=False, # existing project/name ok, do not increment59line_thickness=3, # bounding box thickness (pixels)60hide_labels=False, # hide labels61hide_conf=False, # hide confidences62half=False, # use FP16 half-precision inference63dnn=False, # use OpenCV DNN for ONNX inference64):65source = str(source)66save_img = not nosave and not source.endswith('.txt') # save inference images67is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)68is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))69webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)70if is_url and is_file:71source = check_file(source) # download7273# Directories74save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run75(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir7677# Load model78device = select_device(device)79model = DetectMultiBackend(weights, device=device, dnn=dnn)80stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine81imgsz = check_img_size(imgsz, s=stride) # check image size8283# Half84half &= (pt or jit or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA85if pt or jit:86model.model.half() if half else model.model.float()8788# Dataloader89if webcam:90view_img = check_imshow()91cudnn.benchmark = True # set True to speed up constant image size inference92dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)93bs = len(dataset) # batch_size94else:95dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)96bs = 1 # batch_size97vid_path, vid_writer = [None] * bs, [None] * bs9899# Run inference100model.warmup(imgsz=(1, 3, *imgsz), half=half) # warmup101dt, seen = [0.0, 0.0, 0.0], 0102for path, im, im0s, vid_cap, s in dataset:103t1 = time_sync()104im = torch.from_numpy(im).to(device)105im = im.half() if half else im.float() # uint8 to fp16/32106im /= 255 # 0 - 255 to 0.0 - 1.0107if len(im.shape) == 3:108im = im[None] # expand for batch dim109t2 = time_sync()110dt[0] += t2 - t1111112# Inference113visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False114pred = model(im, augment=augment, visualize=visualize)115t3 = time_sync()116dt[1] += t3 - t2117118# NMS119# pred: list*(n, [xylsθ, conf, cls]) θ ∈ [-pi/2, pi/2)120pred = non_max_suppression_obb(pred, conf_thres, iou_thres, classes, agnostic_nms, multi_label=True, max_det=max_det)121dt[2] += time_sync() - t3122123# Second-stage classifier (optional)124# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)125126# Process predictions127for i, det in enumerate(pred): # per image128pred_poly = rbox2poly(det[:, :5]) # (n, [x1 y1 x2 y2 x3 y3 x4 y4])129seen += 1130if webcam: # batch_size >= 1131p, im0, frame = path[i], im0s[i].copy(), dataset.count132s += f'{i}: '133else:134p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)135136p = Path(p) # to Path137save_path = str(save_dir / p.name) # im.jpg138txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt139s += '%gx%g ' % im.shape[2:] # print string140gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh141imc = im0.copy() if save_crop else im0 # for save_crop142annotator = Annotator(im0, line_width=line_thickness, example=str(names))143if len(det):144# Rescale polys from img_size to im0 size145# det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()146pred_poly = scale_polys(im.shape[2:], pred_poly, im0.shape)147det = torch.cat((pred_poly, det[:, -2:]), dim=1) # (n, [poly conf cls])148149# Print results150for c in det[:, -1].unique():151n = (det[:, -1] == c).sum() # detections per class152s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string153154# Write results155for *poly, conf, cls in reversed(det):156if save_txt: # Write to file157# xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh158poly = poly.tolist()159line = (cls, *poly, conf) if save_conf else (cls, *poly) # label format160with open(txt_path + '.txt', 'a') as f:161f.write(('%g ' * len(line)).rstrip() % line + '\n')162163if save_img or save_crop or view_img: # Add poly to image164c = int(cls) # integer class165label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')166# annotator.box_label(xyxy, label, color=colors(c, True))167annotator.poly_label(poly, label, color=colors(c, True))168if save_crop: # Yolov5-obb doesn't support it yet169# save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)170pass171172# Print time (inference-only)173LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')174175# Stream results176im0 = annotator.result()177if view_img:178cv2.imshow(str(p), im0)179cv2.waitKey(1) # 1 millisecond180181# Save results (image with detections)182if save_img:183if dataset.mode == 'image':184cv2.imwrite(save_path, im0)185else: # 'video' or 'stream'186if vid_path[i] != save_path: # new video187vid_path[i] = save_path188if isinstance(vid_writer[i], cv2.VideoWriter):189vid_writer[i].release() # release previous video writer190if vid_cap: # video191fps = vid_cap.get(cv2.CAP_PROP_FPS)192w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))193h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))194else: # stream195fps, w, h = 30, im0.shape[1], im0.shape[0]196save_path += '.mp4'197vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))198vid_writer[i].write(im0)199200# Print results201t = tuple(x / seen * 1E3 for x in dt) # speeds per image202LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)203if save_txt or save_img:204s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''205LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")206if update:207strip_optimizer(weights) # update model (to fix SourceChangeWarning)208209210def parse_opt():211parser = argparse.ArgumentParser()212parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'runs/train/yolov5n_DroneVehicle/weights/best.pt', help='model path(s)')213parser.add_argument('--source', type=str, default='/media/test/4d846cae-2315-4928-8d1b-ca6d3a61a3c6/DroneVehicle/val/raw/images/', help='file/dir/URL/glob, 0 for webcam')214parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[840], help='inference size h,w')215parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')216parser.add_argument('--iou-thres', type=float, default=0.2, help='NMS IoU threshold')217parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')218parser.add_argument('--device', default='3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')219parser.add_argument('--view-img', action='store_true', help='show results')220parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')221parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')222parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')223parser.add_argument('--nosave', action='store_true', help='do not save images/videos')224parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')225parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')226parser.add_argument('--augment', action='store_true', help='augmented inference')227parser.add_argument('--visualize', action='store_true', help='visualize features')228parser.add_argument('--update', action='store_true', help='update all models')229parser.add_argument('--project', default='runs/detect', help='save results to project/name')230parser.add_argument('--name', default='exp', help='save results to project/name')231parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')232parser.add_argument('--line-thickness', default=2, type=int, help='bounding box thickness (pixels)')233parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')234parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')235parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')236parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')237opt = parser.parse_args()238opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand239print_args(FILE.stem, opt)240return opt241242243def main(opt):244check_requirements(exclude=('tensorboard', 'thop'))245run(**vars(opt))246247248if __name__ == "__main__":249opt = parse_opt()250main(opt)251252253