# focal loss g = h['fl_gamma'] # focal loss gamma if g > 0: BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
# per output nt = 0# targets for i, pi in enumerate(p): # layer index, layer predictions b, a, gj, gi = indices[i] # image, anchor, gridy, gridx tobj = torch.zeros_like(pi[..., 0]) # target obj
nb = b.shape[0] # number of targets if nb: nt += nb # cumulative targets ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
# Obj tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * giou.detach().clamp(0).type(tobj.dtype) # giou ratio
# Class if model.nc > 1: # cls loss (only if multiple classes) t = torch.full_like(ps[:, 5:], cn) # targets t[range(nb), tcls[i]] = cp lcls += BCEcls(ps[:, 5:], t) # BCE
# Append targets to text file # with open('targets.txt', 'a') as file: # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
lobj += BCEobj(pi[..., 4], tobj) # obj loss
lbox *= h['giou'] lobj *= h['obj'] lcls *= h['cls'] bs = tobj.shape[0] # batch size if red == 'sum': g = 3.0# loss gain lobj *= g / bs if nt: lcls *= g / nt / model.nc lbox *= g / nt
loss = lbox + lobj + lcls return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
# Initialize device = torch_utils.select_device(opt.device) if os.path.exists(out): shutil.rmtree(out) # delete output folder os.makedirs(out) # make new output folder half = device.type != 'cpu'# half precision only supported on CUDA
# 下载模型 google_utils.attempt_download(weights) # 加载权重 model = torch.load(weights, map_location=device)['model'].float() # torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning # model.fuse() # 设置模型为推理模式 model.to(device).eval() if half: model.half() # to FP16
# 设置 Dataloader vid_path, vid_writer = None, None if webcam: view_img = True torch.backends.cudnn.benchmark = True# set True to speed up constant image size inference dataset = LoadStreams(source, img_size=imgsz) else: save_img = True dataset = LoadImages(source, img_size=imgsz)
# 获取检测类别的标签名称 names = model.names if hasattr(model, 'names') else model.modules.names # 定义颜色 colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
# 开始推理 t0 = time.time() # 初始化一张全为0的图片 img = torch.zeros((1, 3, imgsz, imgsz), device=device) _ = model(img.half() if half else img) if device.type != 'cpu'elseNone for path, img, im0s, vid_cap in dataset: img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() # uint8 to fp16/32 img /= 255.0# 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0)
# 预测结果 t1 = torch_utils.time_synchronized() pred = model(img, augment=opt.augment)[0]
# Stream results if view_img: cv2.imshow(p, im0) if cv2.waitKey(1) == ord('q'): # q to quit raise StopIteration
# 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)
if save_txt or save_img: print('Results saved to %s' % os.getcwd() + os.sep + out) if platform == 'darwin': # MacOS os.system('open ' + save_path)
原始缩放尺寸是416×416，都除以原始图像的尺寸后，可以得到0.52，和0.69两个缩放系数，选择小的缩放系数0.52。
