(save_img=False)
| 28 | return list(filter(None, names)) # filter removes empty strings (such as last line) |
| 29 | |
| 30 | def detect(save_img=False): |
| 31 | out, source, weights, view_img, save_txt, imgsz, cfg, names = \ |
| 32 | opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, opt.cfg, opt.names |
| 33 | webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt') |
| 34 | |
| 35 | # Initialize |
| 36 | device = select_device(opt.device) |
| 37 | if os.path.exists(out): |
| 38 | shutil.rmtree(out) # delete output folder |
| 39 | os.makedirs(out) # make new output folder |
| 40 | half = device.type != 'cpu' # half precision only supported on CUDA |
| 41 | |
| 42 | # Load model |
| 43 | model = Darknet(cfg, imgsz).cuda() |
| 44 | try: |
| 45 | model.load_state_dict(torch.load(weights[0], map_location=device)['model']) |
| 46 | #model = attempt_load(weights, map_location=device) # load FP32 model |
| 47 | #imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size |
| 48 | except: |
| 49 | load_darknet_weights(model, weights[0]) |
| 50 | model.to(device).eval() |
| 51 | if half: |
| 52 | model.half() # to FP16 |
| 53 | |
| 54 | # Second-stage classifier |
| 55 | classify = False |
| 56 | if classify: |
| 57 | modelc = load_classifier(name='resnet101', n=2) # initialize |
| 58 | modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights |
| 59 | modelc.to(device).eval() |
| 60 | |
| 61 | # Set Dataloader |
| 62 | vid_path, vid_writer = None, None |
| 63 | if webcam: |
| 64 | view_img = True |
| 65 | cudnn.benchmark = True # set True to speed up constant image size inference |
| 66 | dataset = LoadStreams(source, img_size=imgsz) |
| 67 | else: |
| 68 | save_img = True |
| 69 | dataset = LoadImages(source, img_size=imgsz, auto_size=64) |
| 70 | |
| 71 | # Get names and colors |
| 72 | names = load_classes(names) |
| 73 | colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))] |
| 74 | |
| 75 | # Run inference |
| 76 | t0 = time.time() |
| 77 | img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img |
| 78 | _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once |
| 79 | for path, img, im0s, vid_cap in dataset: |
| 80 | img = torch.from_numpy(img).to(device) |
| 81 | img = img.half() if half else img.float() # uint8 to fp16/32 |
| 82 | img /= 255.0 # 0 - 255 to 0.0 - 1.0 |
| 83 | if img.ndimension() == 3: |
| 84 | img = img.unsqueeze(0) |
| 85 | |
| 86 | # Inference |
| 87 | t1 = time_synchronized() |
no test coverage detected