def __init__( self, device="cpu", weights=f"{current_dir}/weights/mobilenet0.25_Final.pth", score_thresh=0.5, top_k=100, nms_thresh=0.4, backbone="mobilenet", use_trt=False, ): assert (device in ["cpu", "cuda"]) or ("cuda" in device) try: device = int(device) device = f"cuda:{device}" except ValueError: pass self.device = torch.device(device) if not torch.cuda.is_available(): self.device = torch.device("cpu") assert backbone in ["mobilenet", "resnet50"] self.cfg = cfg_mnet if backbone == "mobilenet" else cfg_re50 if use_trt and available_trt and (device != "cpu"): if os.path.exists( os.path.join(current_dir, "engines", f"retina_trt_{device}.pth")): print("Load TRT engine...") self.detector = TRTModule() self.detector.load_state_dict( torch.load( os.path.join(current_dir, "engines", f"retina_trt_{device}.pth"))) else: try: self.detector = create_engine(weights, self.device) except Exception as e: print("ERROR: Cannot create engine for retinaface.") print(e) else: self.detector = RetinaFace(self.cfg) load_model(self.detector, weights, self.device) self.detector.eval() self.detector.to(self.device) self.priorbox = PriorBox(self.cfg).forward().to(self.device) self.score_thresh = score_thresh self.top_k = top_k self.nms_thresh = nms_thresh
def load_model(): pretrained_path = 'retinaface/weights/mobilenet0.25_Final.pth' # print('Loading pretrained model from {}'.format(pretrained_path)) model = RetinaFace(cfg=cfg_mnet, phase='test') device = torch.cuda.current_device() pretrained_dict = torch.load( pretrained_path, map_location=lambda storage, loc: storage.cuda(device)) if "state_dict" in pretrained_dict.keys(): pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.') else: pretrained_dict = remove_prefix(pretrained_dict, 'module.') check_keys(model, pretrained_dict) model.load_state_dict(pretrained_dict, strict=False) # print('Finished loading model!') return model
def create_engine(weights, device, eps=1e-3): print("Create trt engine for retintaface...") from torch2trt import torch2trt model = RetinaFace(cfg).to(device) load_model(model, weights, device) model.eval() x = torch.ones((1, 3, cfg["image_size"], cfg["image_size"])).to(device) model_trt = torch2trt(model, [x]) print("Ok. Check outputs...") y = model(x) y_trt = model_trt(x) for out, out_trt in zip(y, y_trt): if torch.max(torch.abs(out - out_trt)) > eps: raise RuntimeError os.makedirs(os.path.join(current_dir, "engines"), exist_ok=True) torch.save(model_trt.state_dict(), os.path.join(current_dir, "engines", f"retina_trt_{device.index}.pth")) return model_trt
def __init__( self, device="cpu", weights=f"{current_dir}/weights/mobilenet0.25_Final.pth", score_thresh=0.5, top_k=100, nms_thresh=0.4, use_trt=False, ): try: device = int(device) cuda_device = f"cuda:{device}" self.device = torch.device(cuda_device) if torch.cuda.is_available() else torch.device("cpu") except ValueError: self.device = torch.device("cpu") if use_trt and available_trt and (device != "cpu"): if os.path.exists(os.path.join(current_dir, "engines", f"retina_trt_{device}.pth")): print("Load TRT engine...") self.detector = TRTModule() self.detector.load_state_dict(torch.load(os.path.join(current_dir, "engines", f"retina_trt_{device}.pth"))) else: try: self.detector = create_engine(weights, self.device) except Exception as e: print("ERROR: Cannot create engine for retinaface.") print(e) else: self.detector = RetinaFace(cfg) load_model(self.detector, weights, self.device) self.detector.eval() self.detector.to(self.device) self.priorbox = PriorBox(cfg).forward().to(self.device) self.score_thresh = score_thresh self.top_k = top_k self.nms_thresh = nms_thresh
def recognise_face(image, feature_type='sift', classifier_type='svm', creative_mode=0): class_names = [ '01', '02', '03', '04', '05', '06', '07', '08', '09', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21', '22', '23', '24', '25', '26', '27', '28', '29', '30', '31', '32', '33', '34', '36', '38', '40', '42', '44', '46', '48', '50', '52', '54', '56', '58', '60', '78' ] # load retinaface model net = RetinaFace(cfg=cfg_re50, phase='test') net = load_model(net, './models/Resnet50_Final.pth', False) net.eval() print('Model loaded successfully') cudnn.benchmark = True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") net = net.to(device) # detect faces in image faces = face_detector.face_detector(image, out_name='results', net=net, save_image=True) img = cv2.imread(image, cv2.IMREAD_COLOR) centres = [] bboxes = [] results = {} if not os.path.exists("./det_faces/"): os.mkdir("./det_faces/") else: files = glob.glob("./det_faces/*") for f in files: os.remove(f) for face in faces: centre = (int((face[2] + face[0]) / 2), int((face[3] + face[1]) / 2)) centres.append(centre) image = img[int(face[1]):int(face[3]), int(face[0]):int(face[2])] cv2.imwrite(f'./det_faces/face_{centre[0]}_{centre[1]}.PNG', image) bbox = [int(f) for f in face[0:4]] bboxes.append(bbox) if creative_mode: whole = cv2.imread('./results.jpg', cv2.IMREAD_COLOR) for face in bboxes: roi = whole[face[1]:face[3], face[0]:face[2], :] cartoon = creative.cartoonify(roi) whole[face[1]:face[3], face[0]:face[2], :] = cartoon cv2.imwrite('results.jpg', whole) if feature_type == 'sift': scaler = joblib.load('./models/SIFT_SCALER.bin') vocab = np.load('./models/sift_vocab.npy') detector = cv2.xfeatures2d.SIFT_create() if feature_type == 'surf': scaler = joblib.load('./models/SURF_SCALER.bin') vocab = np.load('./models/surf_vocab.npy') detector = cv2.xfeatures2d.SURF_create() if classifier_type.lower() == 'cnn': fin, prob = recognise_face_cnn('./det_faces/', class_names) elif classifier_type.lower() == 'svm' or classifier_type.lower() == 'rf': fin = recognise_face_classic('./det_faces/', classifier=classifier_type, extractor=feature_type, detector=detector, scaler=scaler, vocab=vocab) prob = None return fin, prob
cv2.putText(img_raw, text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255)) # save image name = "./" + str(out_name) + ".jpg" cv2.imwrite(name, img_raw) results = [] for face in dets: if face[4] > vis_thres: results.append(face) return results # returns the bounding boxes if __name__ == '__main__': net = RetinaFace(cfg=cfg_re50, phase='test') net = load_model(net, 'Resnet50_Final.pth', False) net.eval() print('Model loaded successfully') cudnn.benchmark = True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") net = net.to(device) individual_list = [ file for file in os.listdir() if 'IndividualImages' in file and ".txt" not in file ] for ind in individual_list: print(f"Working on dir {ind}") im_list = [ file for file in os.listdir(f'./{ind}/') if file.endswith('.JPG')
class RetinaDetector: def __init__( self, device="cpu", weights=f"{current_dir}/weights/mobilenet0.25_Final.pth", score_thresh=0.5, top_k=100, nms_thresh=0.4, backbone="mobilenet", use_trt=False, ): assert (device in ["cpu", "cuda"]) or ("cuda" in device) try: device = int(device) device = f"cuda:{device}" except ValueError: pass self.device = torch.device(device) if not torch.cuda.is_available(): self.device = torch.device("cpu") assert backbone in ["mobilenet", "resnet50"] self.cfg = cfg_mnet if backbone == "mobilenet" else cfg_re50 if use_trt and available_trt and (device != "cpu"): if os.path.exists( os.path.join(current_dir, "engines", f"retina_trt_{device}.pth")): print("Load TRT engine...") self.detector = TRTModule() self.detector.load_state_dict( torch.load( os.path.join(current_dir, "engines", f"retina_trt_{device}.pth"))) else: try: self.detector = create_engine(weights, self.device) except Exception as e: print("ERROR: Cannot create engine for retinaface.") print(e) else: self.detector = RetinaFace(self.cfg) load_model(self.detector, weights, self.device) self.detector.eval() self.detector.to(self.device) self.priorbox = PriorBox(self.cfg).forward().to(self.device) self.score_thresh = score_thresh self.top_k = top_k self.nms_thresh = nms_thresh @staticmethod def aligning(img, lands, crop_size=512): return insightface_align(img, lands, output_size=crop_size) def __call__(self, image): image_ = pad(image) scale = self.cfg["image_size"] / image_.shape[0] image_ = cv2.resize(image_, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR) image_ = image_.astype(np.float32) image_ -= (104, 117, 123) image_ = image_.transpose(2, 0, 1) image_ = torch.from_numpy(image_).unsqueeze(0) image_ = image_.to(self.device) with torch.no_grad(): loc, conf, landms = self.detector(image_) boxes = decode(loc.data.squeeze(0), self.priorbox, self.cfg['variance']) boxes = boxes * image_.size(2) / scale boxes = boxes.cpu().numpy() scores = conf.squeeze(0).cpu().numpy()[:, 1] landms = decode_landm(landms.data.squeeze(0), self.priorbox, self.cfg['variance']) landms = landms * image_.size(2) / scale landms = landms.cpu().numpy() # ignore low scores inds = np.where(scores > self.score_thresh)[0] boxes = boxes[inds] landms = landms[inds] scores = scores[inds] # keep top-K before NMS order = scores.argsort()[::-1][:self.top_k] boxes = boxes[order] landms = landms[order] scores = scores[order] # do NMS dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False) keep = nms(dets, self.nms_thresh) dets = dets[keep, :] landms = landms[keep] boxes = dets[:, :4] scores = dets[:, 4] landms = landms.reshape(-1, 5, 2) return boxes, landms, scores
num_classes = 2 img_dim = cfg['image_size'] num_gpu = cfg['ngpu'] batch_size = cfg['batch_size'] max_epoch = cfg['epoch'] gpu_train = cfg['gpu_train'] num_workers = args.num_workers momentum = args.momentum weight_decay = args.weight_decay initial_lr = args.lr gamma = args.gamma training_dataset = args.training_dataset save_folder = args.save_folder net = RetinaFace(cfg=cfg) print("Printing net...") print(net) if args.resume_net is not None: print('Loading resume network...') state_dict = torch.load(args.resume_net) # create new OrderedDict that does not contain `module.` from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): head = k[:7] if head == 'module.': name = k[7:] # remove `module.` else: name = k