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