jnulzl
commented
1 year ago @rupsgrover Hi, These codes maybe you want:
# ******************************************************
# Author : liuyang && jnulzl
# Last modified : 2022-12-05 15:26
# Filename : demo_single.py
# ******************************************************
from __future__ import absolute_import
import argparse
import time
import numpy as np
import torch
import os
import cv2
from core.workspace import create, load_config
from torch.autograd import Variable
from utils.nms.nms_wrapper import nms
from data import anchor_utils
from data.transform.image_util import normalize_img
from data.datasets_utils.visualize import draw_bboxes
from data.anchors_opr.generate_anchors import GeneartePriorBoxes
parser = argparse.ArgumentParser(description='Test Details')
parser.add_argument('--weight_path', default='snapshots/MogFace_Ali-AMS/model_70000.pth',
type=str, help='The weight path.')
parser.add_argument('--source', default='', type=str, help='img path')
parser.add_argument('--nms_th', default=0.3, type=float, help='nms threshold.')
parser.add_argument('--pre_nms_top_k', default=5000, type=int, help='number of max score image.')
parser.add_argument('--score_th', default=0.01, type=float, help='score threshold.')
parser.add_argument('--max_bbox_per_img', default=750, type=int, help='max number of det bbox.')
parser.add_argument('--config', '-c', default='./config.yml', type=str, help='config yml.')
parser.add_argument('--test_idx', default=None, type=int)
class DataSetting(object):
def __init__(self):
pass
def detect_face(net, image, shrink, generate_anchors_fn):
x = image
if shrink != 1:
x = cv2.resize(image, None, None, fx=shrink, fy=shrink, interpolation=cv2.INTER_LINEAR)
# print('shrink:{}'.format(shrink))
width = x.shape[1]
height = x.shape[0]
print('width: {}, height: {}'.format(width, height))
x = torch.from_numpy(x).permute(2, 0, 1)
x = x.unsqueeze(0)
x = Variable(x.cuda(), volatile=True)
out = net(x)
anchors = anchor_utils.transform_anchor((generate_anchors_fn(height, width)))
anchors = torch.FloatTensor(anchors).cuda()
decode_bbox = anchor_utils.decode(out[1].squeeze(0), anchors)
boxes = decode_bbox
scores = out[0].squeeze(0)
select_idx_list = []
tmp_height = height
tmp_width = width
test_idx = args.test_idx
if test_idx is not None:
for i in range(2):
tmp_height = (tmp_height + 1) // 2
tmp_width = (tmp_width + 1) // 2
for i in range(6):
if i == 0:
select_idx_list.append(tmp_height * tmp_width)
else:
select_idx_list.append(tmp_height * tmp_width + select_idx_list[i - 1])
tmp_height = (tmp_height + 1) // 2
tmp_width = (tmp_width + 1) // 2
if test_idx == 2:
boxes = boxes[:select_idx_list[(test_idx - 2)]]
scores = scores[:select_idx_list[(test_idx - 2)]]
else:
boxes = boxes[select_idx_list[test_idx - 3]: select_idx_list[test_idx - 2]]
scores = scores[select_idx_list[test_idx - 3]: select_idx_list[test_idx - 2]]
# print('scores shape', scores.shape)
# print('boxes shape', boxes.shape)
top_k = args.pre_nms_top_k
v, idx = scores[:, 0].sort(0)
idx = idx[-top_k:]
boxes = boxes[idx]
scores = scores[idx]
# [11620, 4]
boxes = boxes.cpu().numpy()
w = boxes[:, 2] - boxes[:, 0] + 1
h = boxes[:, 3] - boxes[:, 1] + 1
boxes[:, 0] /= shrink
boxes[:, 1] /= shrink
boxes[:, 2] = boxes[:, 0] + w / shrink - 1
boxes[:, 3] = boxes[:, 1] + h / shrink - 1
# boxes = boxes / shrink
# [11620, 2]
scores = scores.cpu().numpy()
inds = np.where(scores[:, 0] > args.score_th)[0]
if len(inds) == 0:
det = np.empty([0, 5], dtype=np.float32)
return det
c_bboxes = boxes[inds]
# [5,]
c_scores = scores[inds, 0]
# [5, 5]
c_dets = np.hstack((c_bboxes, c_scores[:, np.newaxis])).astype(np.float32, copy=False)
keep = nms(c_dets, args.nms_th)
c_dets = c_dets[keep, :]
max_bbox_per_img = args.max_bbox_per_img
if max_bbox_per_img > 0:
image_scores = c_dets[:, -1]
if len(image_scores) > max_bbox_per_img:
image_thresh = np.sort(image_scores)[-max_bbox_per_img]
keep = np.where(c_dets[:, -1] >= image_thresh)[0]
c_dets = c_dets[keep, :]
for i in range(c_dets.shape[0]):
if c_dets[i][0] < 0.0:
c_dets[i][0] = 0.0
if c_dets[i][1] < 0.0:
c_dets[i][1] = 0.0
if c_dets[i][2] > width - 1:
c_dets[i][2] = width - 1
if c_dets[i][3] > height - 1:
c_dets[i][3] = height - 1
return c_dets
def process_img(img, net, generate_anchors_fn, normalize_setting):
with torch.no_grad():
img = normalize_img(img.astype(np.float32), normalize_setting)
max_im_shrink = (0x7fffffff / 200.0 / (
img.shape[0] * img.shape[1])) ** 0.5 # the max size of input image for caffe
max_im_shrink = 2.2 if max_im_shrink > 2.2 else max_im_shrink
shrink = max_im_shrink if max_im_shrink < 1 else 1
boxes = detect_face(net, img, shrink, generate_anchors_fn) # origin test
return boxes
if __name__ == "__main__":
args = parser.parse_args()
cfg = load_config(args.config)
generate_anchors_fn = GeneartePriorBoxes(
scale_list=cfg['GeneartePriorBoxes']['scale_list'],
aspect_ratio_list=cfg['GeneartePriorBoxes']['aspect_ratio_list'],
stride_list=cfg['GeneartePriorBoxes']['stride_list'],
anchor_size_list=cfg['GeneartePriorBoxes']['anchor_size_list'],
)
normalize_setting = DataSetting()
normalize_setting.use_rgb = cfg['BasePreprocess']['use_rgb']
normalize_setting.img_mean = (np.array(cfg['BasePreprocess']['img_mean']).astype('float32') * 255)[::-1]
normalize_setting.img_std = (np.array(cfg['BasePreprocess']['img_std']).astype('float32') * 255)[::-1]
normalize_setting.normalize_pixel = cfg['BasePreprocess']['normalize_pixel']
net = create(cfg.architecture)
print('Load model from {}'.format(args.weight_path))
net.load_state_dict(torch.load(args.weight_path))
net.cuda()
net.eval()
onnx_path = args.weight_path.replace('.pth', '.onnx')
if not os.path.exists(onnx_path):
img = torch.rand(1, 3, 640, 640)
model = net.cpu()
torch.onnx.export(model, img, onnx_path, verbose=True, opset_version=11,
input_names=['input'],
output_names=['output1', 'output2']
)
print('Finish load model.')
# For test image list
if args.source.endswith(".txt"):
with open(args.source, 'r') as fpR:
lines = fpR.readlines()
for line in lines:
img = cv2.imread(line.strip())
img_show = img.copy()
t1 = time.time()
boxes = process_img(img, net, generate_anchors_fn, normalize_setting)
t2 = time.time()
print("Inference time : %d" % ((t2 - t1) * 1000))
img = draw_bboxes(img_show, boxes[:, :4])
cv2.imshow("Demo", img)
if 27 == cv2.waitKey(0):
break
# For test video
elif args.source.endswith(".mp4") or args.source.endswith(".avi"):
cap = cv2.VideoCapture(args.source)
while True:
ret, img = cap.read()
if img is None:
break
img_show = img.copy()
t1 = time.time()
boxes = process_img(img, net, generate_anchors_fn, normalize_setting)
t2 = time.time()
print("Inference time : %d" % ((t2 - t1) * 1000))
img = draw_bboxes(img_show, boxes[:, :4])
cv2.imshow("Demo", img)
if 27 == cv2.waitKey(1):
break
# For test single image
else:
with torch.no_grad():
img = cv2.imread(args.source)
img_show = img.copy()
t1 = time.time()
boxes = process_img(img, net, generate_anchors_fn, normalize_setting)
t2 = time.time()
print("Inference time : %d" % ((t2 - t1) * 1000))
img = draw_bboxes(img_show, boxes[:, :4])
cv2.imwrite("Demo.jpg", img)
cv2.imshow("Demo", img)
cv2.waitKey(0)Usage:
# MogFace_Ali-AMS url : https://drive.google.com/drive/folders/1wMrfwrYSirGnLcMUC_KCCCr0Xd12uVzk?usp=sharing
CUDA_VISIBLE_DEVICES=0 python demo_single.py --config configs/mogface/MogFace_Ali-AMS.yml --weight_path snapshots/MogFace_Ali-AMS/model_70000.pth --source dataset/video3.mp4
There is no test/detect_face function for detecting faces in images other than the widerface dataset.