|
| 1 | +# YOLOv5 🚀 by Ultralytics, GPL-3.0 license |
| 2 | +""" |
| 3 | +Run YOLOv5 classification inference on images, videos, directories, globs, YouTube, webcam, streams, etc. |
| 4 | +
|
| 5 | +Usage - sources: |
| 6 | + $ python classify/predict.py --weights yolov5s-cls.pt --source 0 # webcam |
| 7 | + img.jpg # image |
| 8 | + vid.mp4 # video |
| 9 | + screen # screenshot |
| 10 | + path/ # directory |
| 11 | + 'path/*.jpg' # glob |
| 12 | + 'https://youtu.be/Zgi9g1ksQHc' # YouTube |
| 13 | + 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream |
| 14 | +
|
| 15 | +Usage - formats: |
| 16 | + $ python classify/predict.py --weights yolov5s-cls.pt # PyTorch |
| 17 | + yolov5s-cls.torchscript # TorchScript |
| 18 | + yolov5s-cls.onnx # ONNX Runtime or OpenCV DNN with --dnn |
| 19 | + yolov5s-cls_openvino_model # OpenVINO |
| 20 | + yolov5s-cls.engine # TensorRT |
| 21 | + yolov5s-cls.mlmodel # CoreML (macOS-only) |
| 22 | + yolov5s-cls_saved_model # TensorFlow SavedModel |
| 23 | + yolov5s-cls.pb # TensorFlow GraphDef |
| 24 | + yolov5s-cls.tflite # TensorFlow Lite |
| 25 | + yolov5s-cls_edgetpu.tflite # TensorFlow Edge TPU |
| 26 | + yolov5s-cls_paddle_model # PaddlePaddle |
| 27 | +""" |
| 28 | + |
| 29 | +import argparse |
| 30 | +import os |
| 31 | +import platform |
| 32 | +import sys |
| 33 | +from pathlib import Path |
| 34 | + |
| 35 | +import torch |
| 36 | +import torch.nn.functional as F |
| 37 | + |
| 38 | +FILE = Path(__file__).resolve() |
| 39 | +ROOT = FILE.parents[1] # YOLOv5 root directory |
| 40 | +if str(ROOT) not in sys.path: |
| 41 | + sys.path.append(str(ROOT)) # add ROOT to PATH |
| 42 | +ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative |
| 43 | + |
| 44 | +from models.common import DetectMultiBackend |
| 45 | +from utils.augmentations import classify_transforms |
| 46 | +from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams |
| 47 | +from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2, |
| 48 | + increment_path, print_args, strip_optimizer) |
| 49 | +from utils.plots import Annotator |
| 50 | +from utils.torch_utils import select_device, smart_inference_mode |
| 51 | + |
| 52 | + |
| 53 | +@smart_inference_mode() |
| 54 | +def run( |
| 55 | + weights=ROOT / 'yolov5s-cls.pt', # model.pt path(s) |
| 56 | + source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam) |
| 57 | + data=ROOT / 'data/coco128.yaml', # dataset.yaml path |
| 58 | + imgsz=(224, 224), # inference size (height, width) |
| 59 | + device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu |
| 60 | + view_img=False, # show results |
| 61 | + save_txt=False, # save results to *.txt |
| 62 | + nosave=False, # do not save images/videos |
| 63 | + augment=False, # augmented inference |
| 64 | + visualize=False, # visualize features |
| 65 | + update=False, # update all models |
| 66 | + project=ROOT / 'runs/predict-cls', # save results to project/name |
| 67 | + name='exp', # save results to project/name |
| 68 | + exist_ok=False, # existing project/name ok, do not increment |
| 69 | + half=False, # use FP16 half-precision inference |
| 70 | + dnn=False, # use OpenCV DNN for ONNX inference |
| 71 | + vid_stride=1, # video frame-rate stride |
| 72 | +): |
| 73 | + source = str(source) |
| 74 | + save_img = not nosave and not source.endswith('.txt') # save inference images |
| 75 | + is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) |
| 76 | + is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')) |
| 77 | + webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file) |
| 78 | + screenshot = source.lower().startswith('screen') |
| 79 | + if is_url and is_file: |
| 80 | + source = check_file(source) # download |
| 81 | + |
| 82 | + # Directories |
| 83 | + save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run |
| 84 | + (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir |
| 85 | + |
| 86 | + # Load model |
| 87 | + device = select_device(device) |
| 88 | + model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half) |
| 89 | + stride, names, pt = model.stride, model.names, model.pt |
| 90 | + imgsz = check_img_size(imgsz, s=stride) # check image size |
| 91 | + |
| 92 | + # Dataloader |
| 93 | + bs = 1 # batch_size |
| 94 | + if webcam: |
| 95 | + view_img = check_imshow(warn=True) |
| 96 | + dataset = LoadStreams(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride) |
| 97 | + bs = len(dataset) |
| 98 | + elif screenshot: |
| 99 | + dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt) |
| 100 | + else: |
| 101 | + dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride) |
| 102 | + vid_path, vid_writer = [None] * bs, [None] * bs |
| 103 | + |
| 104 | + # Run inference |
| 105 | + model.warmup(imgsz=(1 if pt else bs, 3, *imgsz)) # warmup |
| 106 | + seen, windows, dt = 0, [], (Profile(), Profile(), Profile()) |
| 107 | + for path, im, im0s, vid_cap, s in dataset: |
| 108 | + with dt[0]: |
| 109 | + im = torch.Tensor(im).to(model.device) |
| 110 | + im = im.half() if model.fp16 else im.float() # uint8 to fp16/32 |
| 111 | + if len(im.shape) == 3: |
| 112 | + im = im[None] # expand for batch dim |
| 113 | + |
| 114 | + # Inference |
| 115 | + with dt[1]: |
| 116 | + results = model(im) |
| 117 | + |
| 118 | + # Post-process |
| 119 | + with dt[2]: |
| 120 | + pred = F.softmax(results, dim=1) # probabilities |
| 121 | + |
| 122 | + # Process predictions |
| 123 | + for i, prob in enumerate(pred): # per image |
| 124 | + seen += 1 |
| 125 | + if webcam: # batch_size >= 1 |
| 126 | + p, im0, frame = path[i], im0s[i].copy(), dataset.count |
| 127 | + s += f'{i}: ' |
| 128 | + else: |
| 129 | + p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0) |
| 130 | + |
| 131 | + p = Path(p) # to Path |
| 132 | + save_path = str(save_dir / p.name) # im.jpg |
| 133 | + txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt |
| 134 | + |
| 135 | + s += '%gx%g ' % im.shape[2:] # print string |
| 136 | + annotator = Annotator(im0, example=str(names), pil=True) |
| 137 | + |
| 138 | + # Print results |
| 139 | + top5i = prob.argsort(0, descending=True)[:5].tolist() # top 5 indices |
| 140 | + s += f"{', '.join(f'{names[j]} {prob[j]:.2f}' for j in top5i)}, " |
| 141 | + |
| 142 | + # Write results |
| 143 | + text = '\n'.join(f'{prob[j]:.2f} {names[j]}' for j in top5i) |
| 144 | + if save_img or view_img: # Add bbox to image |
| 145 | + annotator.text((32, 32), text, txt_color=(255, 255, 255)) |
| 146 | + if save_txt: # Write to file |
| 147 | + with open(f'{txt_path}.txt', 'a') as f: |
| 148 | + f.write(text + '\n') |
| 149 | + |
| 150 | + # Stream results |
| 151 | + im0 = annotator.result() |
| 152 | + if view_img: |
| 153 | + if platform.system() == 'Linux' and p not in windows: |
| 154 | + windows.append(p) |
| 155 | + cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) # allow window resize (Linux) |
| 156 | + cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0]) |
| 157 | + cv2.imshow(str(p), im0) |
| 158 | + cv2.waitKey(1) # 1 millisecond |
| 159 | + |
| 160 | + # Save results (image with detections) |
| 161 | + if save_img: |
| 162 | + if dataset.mode == 'image': |
| 163 | + cv2.imwrite(save_path, im0) |
| 164 | + else: # 'video' or 'stream' |
| 165 | + if vid_path[i] != save_path: # new video |
| 166 | + vid_path[i] = save_path |
| 167 | + if isinstance(vid_writer[i], cv2.VideoWriter): |
| 168 | + vid_writer[i].release() # release previous video writer |
| 169 | + if vid_cap: # video |
| 170 | + fps = vid_cap.get(cv2.CAP_PROP_FPS) |
| 171 | + w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| 172 | + h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| 173 | + else: # stream |
| 174 | + fps, w, h = 30, im0.shape[1], im0.shape[0] |
| 175 | + save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos |
| 176 | + vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) |
| 177 | + vid_writer[i].write(im0) |
| 178 | + |
| 179 | + # Print time (inference-only) |
| 180 | + LOGGER.info(f"{s}{dt[1].dt * 1E3:.1f}ms") |
| 181 | + |
| 182 | + # Print results |
| 183 | + t = tuple(x.t / seen * 1E3 for x in dt) # speeds per image |
| 184 | + LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t) |
| 185 | + if save_txt or save_img: |
| 186 | + s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' |
| 187 | + LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}") |
| 188 | + if update: |
| 189 | + strip_optimizer(weights[0]) # update model (to fix SourceChangeWarning) |
| 190 | + |
| 191 | + |
| 192 | +def parse_opt(): |
| 193 | + parser = argparse.ArgumentParser() |
| 194 | + parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model path(s)') |
| 195 | + parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)') |
| 196 | + parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path') |
| 197 | + parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[224], help='inference size h,w') |
| 198 | + parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') |
| 199 | + parser.add_argument('--view-img', action='store_true', help='show results') |
| 200 | + parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') |
| 201 | + parser.add_argument('--nosave', action='store_true', help='do not save images/videos') |
| 202 | + parser.add_argument('--augment', action='store_true', help='augmented inference') |
| 203 | + parser.add_argument('--visualize', action='store_true', help='visualize features') |
| 204 | + parser.add_argument('--update', action='store_true', help='update all models') |
| 205 | + parser.add_argument('--project', default=ROOT / 'runs/predict-cls', help='save results to project/name') |
| 206 | + parser.add_argument('--name', default='exp', help='save results to project/name') |
| 207 | + parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') |
| 208 | + parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') |
| 209 | + parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference') |
| 210 | + parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride') |
| 211 | + opt = parser.parse_args() |
| 212 | + opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand |
| 213 | + print_args(vars(opt)) |
| 214 | + return opt |
| 215 | + |
| 216 | + |
| 217 | +def main(opt): |
| 218 | + check_requirements(exclude=('tensorboard', 'thop')) |
| 219 | + run(**vars(opt)) |
| 220 | + |
| 221 | + |
| 222 | +if __name__ == "__main__": |
| 223 | + opt = parse_opt() |
| 224 | + main(opt) |
0 commit comments