def train(net): net.train() priorbox = PriorBox() with torch.no_grad(): priors = priorbox.forward() priors = priors.to(device) dataloader = DataLoader(VOCDetection(), batch_size=2, collate_fn=detection_collate, num_workers=12) for epoch in range(1000): loss_ls, loss_cs = [], [] load_t0 = time.time() if epoch > 500: adjust_learning_rate(optimizer, 1e-4) for images, targets in dataloader: images = images.to(device) targets = [anno.to(device) for anno in targets] out = net(images) optimizer.zero_grad() loss_l, loss_c = criterion(out, priors, targets) loss = 2 * loss_l + loss_c loss.backward() optimizer.step() loss_cs.append(loss_c.item()) loss_ls.append(loss_l.item()) load_t1 = time.time() print(f'{np.mean(loss_cs)}, {np.mean(loss_ls)} time:{load_t1-load_t0}') torch.save(net.state_dict(), 'Final_FaceBoxes.pth')
def train(net): net.train() priorbox = PriorBox() with torch.no_grad(): priors = priorbox.forward() priors = priors.to(device) dataloader = DataLoader(VOCDetection(), batch_size=2, collate_fn=detection_collate, num_workers=12) for epoch in range(1000): loss_ls, loss_cs = [], [] load_t0 = time.time() if epoch > 500: adjust_learning_rate(optimizer, 1e-4) for images, targets in dataloader: images = images.to(device) targets = [anno.to(device) for anno in targets] out = net(images) optimizer.zero_grad() loss_l, loss_c = criterion(out, priors, targets) loss = 2 * loss_l + loss_c loss.backward() optimizer.step() loss_cs.append(loss_c.item()) loss_ls.append(loss_l.item()) load_t1 = time.time() print(f'{np.mean(loss_cs)}, {np.mean(loss_ls)} time:{load_t1-load_t0}') torch.save(net.state_dict(), 'Final_FaceBoxes.pth')
def init_priors(self, feature_maps, image_size): # Hacky key system, but works.... key = ".".join([str(item) for i in range(len(feature_maps)) for item in feature_maps[i]]) + \ "," + ".".join([str(_) for _ in image_size]) if key in self.prior_cache: return self.prior_cache[key].clone() priorbox = PriorBox(self.cfg, image_size, feature_maps) prior = priorbox.forward() self.prior_cache[key] = prior.clone() return prior
def __init__(self, phase, size, Backbone, Neck, Head, cfg): super(SSD, self).__init__() self.phase = phase self.cfg = cfg self.priorbox = PriorBox(self.cfg) self.priors = self.priorbox.forward() self.size = size # SSD network self.backbone = Backbone self.neck = Neck self.head = Head self.num_classes = cfg['num_classes'] self.softmax = nn.Softmax(dim=-1) self.detect = Detect(self.num_classes , 0, 200, 0.01, 0.45,variance = cfg['variance'], nms_kind=cfg['nms_kind'], beta1=cfg['beta1'])
def __init__(self, model_path, gpu_ids, layers, score_thresh=0.5): """ 检测整体基本流程 :param model_path: 模型路径 :param gpu_ids: gpu序列号 :param layers: 18 , 50 :param score_thresh: 置信度过滤 """ self.keep_top_k = 100 self.nms_threshold = 0.3 self.nms_score = score_thresh self.nms_threshold = self.nms_threshold self.test_size = 640 self.__model_path = model_path self.__gpu_ids = gpu_ids self.device = torch.device('cuda:{}'.format(str(gpu_ids)) if torch. cuda.is_available() else 'cpu') self.layers = layers self.model = RetinaFace(self.layers) self.model = self.__load_model(self.model, self.__model_path) self.model = self.model.to(self.device) self.model.eval() self.priorbox = PriorBox(box_specs_list=[[(0.6, 0.5), (0.75, 1.), (0.9, 1.)], [(0.2, 0.5), (0.4, 1.), (0.6, 1.)], [(0.05, 0.5), (0.1, 1.), (0.2, 1.)], [(0.0125, 0.5), (0.025, 1.), (0.05, 1.)]], base_anchor_size=[1.0, 1.0]) self.priors = self.priorbox.generate(feature_map_shape_list=[(10, 10), (20, 20), (40, 40), (80, 80)], im_height=640, im_width=640) self.priors = self.priors.to(self.device) self.mean = torch.Tensor([104, 117, 123]).to(self.device) self.variance = torch.Tensor([0.1, 0.2]).to(self.device) self.Decode = Decode(self.priors.data, self.variance)
def __init__(self, confidence_threshold=0.02, top_k=1000, nms_threshold=0.4, keep_top_k=500, vis_thres=0.6): self.net = Retina(cfg=cfg_mnet).to(device) self.net = load_model(self.net, 'mnet_plate.pth', False) self.net.eval() self.lprnet = plate_recogition() self.priorbox = PriorBox(cfg_mnet) self.confidence_threshold = confidence_threshold self.top_k = top_k self.nms_threshold = nms_threshold self.keep_top_k = keep_top_k self.vis_thres = vis_thres self.resize = 1 self.points_ref = np.float32([[0, 0], [94, 0], [0, 24], [94, 24]])
class SSD(nn.Module): """Single Shot Multibox Architecture The network is composed of a base VGG network followed by the added multibox conv layers. Each multibox layer branches into 1) conv2d for class conf scores 2) conv2d for localization predictions 3) associated priorbox layer to produce default bounding boxes specific to the layer's feature map size. See: https://arxiv.org/pdf/1512.02325.pdf for more details. Args: phase: (string) Can be "test" or "train" size: input image size base: VGG16 layers for input, size of either 300 or 500 extras: extra layers that feed to multibox loc and conf layers head: "multibox head" consists of loc and conf conv layers """ def __init__(self, phase, size, Basenet, Neck, Head, cfg): super(SSD, self).__init__() self.phase = phase self.cfg = cfg self.priorbox = PriorBox(self.cfg) self.priors = self.priorbox.forward() self.size = size # SSD network self.basenet = Basenet self.neck = Neck self.head = Head self.num_classes = cfg['num_classes'] self.softmax = nn.Softmax(dim=-1) self.detect = Detect(self.num_classes, 0, 200, 0.01, 0.45, variance=cfg['variance'], nms_kind=cfg['nms_kind'], beta1=cfg['beta1']) def forward(self, x, phase): """Applies network layers and ops on input image(s) x. Args: x: input image or batch of images. Shape: [batch,3,300,300]. Return: Depending on phase: test: Variable(tensor) of output class label predictions, confidence score, and corresponding location predictions for each object detected. Shape: [batch,topk,7] train: list of concat outputs from: 1: confidence layers, Shape: [batch*num_priors,num_classes] 2: localization layers, Shape: [batch,num_priors*4] 3: priorbox layers, Shape: [2,num_priors*4] """ x = self.basenet(x) if self.neck is not None: x = self.neck(x) conf, loc = self.head(x) loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1) conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1) if phase == "test": output = self.detect.trans( loc.view(loc.size(0), -1, 4), # loc preds self.softmax(conf.view(conf.size(0), -1, self.num_classes)), # conf preds #self.priors.type(type(x.data)) # default boxes self.priors ) else: output = ( loc.view(loc.size(0), -1, 4), conf.view(conf.size(0), -1, self.num_classes), self.priors ) return output def load_weights(self, base_file): other, ext = os.path.splitext(base_file) if ext == '.pkl' or '.pth': print('Loading weights into state dict...') self.load_state_dict(torch.load(base_file, map_location=lambda storage, loc: storage)) print('Finished!') else: print('Sorry only .pth and .pkl files supported.')
print('Loading Dataset...') (show_classes, num_classes, dataset, epoch_size, max_iter, testset) = load_dataset() print('Loading Network...') from models.detector import Detector model = Detector(args.size, num_classes, args.backbone, args.neck) model.train() model.cuda() num_param = sum(p.numel() for p in model.parameters() if p.requires_grad) print('Total param is : {:e}'.format(num_param)) print('Preparing Optimizer & AnchorBoxes...') optimizer = optim.SGD(tencent_trick(model), lr=args.lr, momentum=0.9, weight_decay=0.0005) criterion = MultiBoxLoss(num_classes, mutual_guide=args.mutual_guide) priorbox = PriorBox(args.base_anchor_size, args.size) with torch.no_grad(): priors = priorbox.forward() priors = priors.cuda() if args.trained_model is not None: print('loading weights from', args.trained_model) state_dict = torch.load(args.trained_model) model.load_state_dict(state_dict, strict=True) else: print('Training {}-{} on {} with {} images'.format(args.neck, args.backbone, dataset.name, len(dataset))) os.makedirs(args.save_folder, exist_ok=True) epoch = 0 timer = Timer() for iteration in range(max_iter): if iteration % epoch_size == 0:
class FaceDetector(object): def __init__(self, model_path, gpu_ids, layers, score_thresh=0.5): """ 检测整体基本流程 :param model_path: 模型路径 :param gpu_ids: gpu序列号 :param layers: 18 , 50 :param score_thresh: 置信度过滤 """ self.keep_top_k = 100 self.nms_threshold = 0.3 self.nms_score = score_thresh self.nms_threshold = self.nms_threshold self.test_size = 640 self.__model_path = model_path self.__gpu_ids = gpu_ids self.device = torch.device('cuda:{}'.format(str(gpu_ids)) if torch. cuda.is_available() else 'cpu') self.layers = layers self.model = RetinaFace(self.layers) self.model = self.__load_model(self.model, self.__model_path) self.model = self.model.to(self.device) self.model.eval() self.priorbox = PriorBox(box_specs_list=[[(0.6, 0.5), (0.75, 1.), (0.9, 1.)], [(0.2, 0.5), (0.4, 1.), (0.6, 1.)], [(0.05, 0.5), (0.1, 1.), (0.2, 1.)], [(0.0125, 0.5), (0.025, 1.), (0.05, 1.)]], base_anchor_size=[1.0, 1.0]) self.priors = self.priorbox.generate(feature_map_shape_list=[(10, 10), (20, 20), (40, 40), (80, 80)], im_height=640, im_width=640) self.priors = self.priors.to(self.device) self.mean = torch.Tensor([104, 117, 123]).to(self.device) self.variance = torch.Tensor([0.1, 0.2]).to(self.device) self.Decode = Decode(self.priors.data, self.variance) def detect(self, detect_input): images, percent = self.preprocess(detect_input) #前处理 loc, conf, landms = self.inference(images) #推理 boxes, landms, scores = self.decode(loc, conf, landms, percent) #解码 boxes, landms, scores = self.postprocess(boxes, landms, scores) #后处理 if len(scores) == 0: return None, None, None else: return boxes, landms, scores def __check_keys(self, model, pretrained_state_dict): ckpt_keys = set(pretrained_state_dict.keys()) model_keys = set(model.state_dict().keys()) used_pretrained_keys = model_keys & ckpt_keys unused_pretrained_keys = ckpt_keys - model_keys missing_keys = model_keys - ckpt_keys print('Missing keys:{}'.format(len(missing_keys))) print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys))) print('Used keys:{}'.format(len(used_pretrained_keys))) assert len( used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint' return True def __remove_prefix(self, state_dict, prefix): ''' Old style model is stored with all names of parameters sharing common prefix 'module.' ''' print('remove prefix \'{}\''.format(prefix)) f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x return {f(key): value for key, value in state_dict.items()} def __load_model(self, model, model_path): print('Loading pretrained model from {}'.format(model_path)) if self.__gpu_ids == None: device = torch.cuda.current_device() pretrained_dict = torch.load( model_path, map_location=lambda storage, loc: storage.cuda(device)) else: pretrained_dict = torch.load( model_path, map_location=lambda storage, loc: storage) if "state_dict" in pretrained_dict.keys(): pretrained_dict = self.__remove_prefix( pretrained_dict['state_dict'], 'module.') else: pretrained_dict = self.__remove_prefix(pretrained_dict, 'module.') self.__check_keys(model, pretrained_dict) model.load_state_dict(pretrained_dict, strict=False) return model def preprocess(self, image): """Preprocess""" image_t, percent = no_deform_resize_pad(image, self.test_size) image_t = rgb_mean_gpu(image_t, self.mean, self.device) images = image_t.permute(2, 0, 1).reshape(1, 3, self.test_size, self.test_size) return images, percent def inference(self, images): """网络推理""" with torch.no_grad(): loc, conf, landms = self.model(images) return loc, conf, landms def decode(self, loc, conf, landms, percent): """推理结果的解码""" boxes = self.Decode.decode_bbox(loc.squeeze(0).data) landms = self.Decode.decode_landm(landms.squeeze(0).data) detect_boxes = boxes * self.test_size * percent detect_landmas = landms * self.test_size * percent scores = conf.squeeze(0).data[:, 1] return detect_boxes, detect_landmas, scores def postprocess(self, boxes, landms, scores): """后处理NMS""" inds = torch.where(scores >= self.nms_score)[0] boxes = boxes[inds] scores = scores[inds] landms = landms[inds] keep = nms(boxes, scores, self.nms_threshold) boxes = boxes[keep] scores = scores[keep] landms = landms[keep] boxes = boxes[:self.keep_top_k, :].cpu().numpy() scores = scores[:self.keep_top_k].cpu().numpy() landms = landms[:self.keep_top_k, :].cpu().numpy() return boxes, landms, scores
"output3": {0: "batch_size"}} torch.onnx.export(net, dummy_input, onnx_output, verbose=True, input_names=input_names, output_names=output_names, opset_version=12, dynamic_axes=dynamic_axes) if False: model = onnx.load(onnx_output) model_simp, check = simplify(model) assert check, "Simplified ONNX model could not be validated" output_path = 'simp.onnx' onnx.save(model_simp, output_path) print('finished exporting onnx ') img_path = 'export/028125-87_110-204&496_524&585-506&564_204&585_210&514_524&496-0_0_5_24_29_33_24_24-52-45.jpg' priorbox = PriorBox(cfg_mnet) points_ref = np.float32([[0, 0], [94, 0], [0, 24], [94, 24]]) confidence_threshold=0.02 top_k=1000 nms_threshold=0.4 keep_top_k=500 vis_thres=0.6 srcimg = cv2.imread(img_path) img = srcimg.astype('float32') im_height, im_width, _ = img.shape img -= (104, 117, 123) with torch.no_grad(): scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]).to(device) img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(device) loc, conf, landms = net(img) # forward pass