def net_build():
ssd_net = build_sfd('test', 640, 2)
net = ssd_net
net.load_state_dict(torch.load(args.resume)) # load model parameters only
net.cuda()
net.eval() # set model as evaluations only dropout and BatchNorm exsist
print('Finished loading model!')
import torchvision.transforms as transforms
from torch.autograd import Variable
import scipy.io as sio
from PIL import Image, ImageDraw
from pyramid import build_sfd
from layers import *
import cv2
import numpy as np
import math
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
torch.cuda.set_device(0)
print('Loading model..')
ssd_net = build_sfd('test', 640, 2)
net = ssd_net
net.load_state_dict(torch.load('./Res50_pyramid.pth'))
net.cuda()
net.eval()
print('Finished loading model!')
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir',
default='ai_challenger_keypoint_test_a_20180103')
parser.add_argument('--out_dir', default='annotations')
parser.add_argument('--image_dir', default='keypoint_test_a_images_20180103')
parser.add_argument('--json_file',
default='keypoint_test_a_annotations_20180103.json')
parser.add_argument('--confidence', type=float, default=0.1)
args = parser.parse_args()
iter_size = accum_batch_size / batch_size
# max_iter = 120000
max_iter = 30600
weight_decay = 0.0001
# stepvalues = (80000, 100000, 120000)
stepvalues = (0, 400, 800, 1200, 1600, 10600, 16600, 22600)
gamma = 0.1
# momentum = 0.9
momentum = 0.99
if args.visdom:
import visdom
viz = visdom.Visdom()
ssd_net = build_sfd('train', 640, num_classes)
net = ssd_net
if args.cuda:
net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True
if args.cuda:
net = net.cuda()
def xavier(param):
init.xavier_uniform(param)
def weights_init(m):
import math
import argparse
import datetime
parser = argparse.ArgumentParser(description='Single Shot MultiBox Detector Training')
# 行人数量在150~220之间使用阈值0.2,高于350使用0.15
parser.add_argument('--probability', default='0.23') # 0.15, 0.2, 0.25，0.3
parser.add_argument('--resume', default='weights/best_our_Res50_pyramid_aug.pth')
parser.add_argument('--gpu', default='0')
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
torch.cuda.set_device(0)
print('Loading model..')
ssd_net = build_sfd('test', 160, 2)
net = ssd_net
net.load_state_dict(torch.load(args.resume)) # load model parameters only
net.cuda()
net.eval() # set model as evaluations only dropout and BatchNorm exsist
print('Finished loading model!')
# global variables
# method 1
def detect_face(image, shrink):
# print('image_shape = ',image.shape) # image_shape = (2160,3840) ->(1080,1920)
# print('shrink_0 = ',shrink) # shrink = 1
x = image
if shrink != 1:
x = cv2.resize(image, None, None, fx=shrink, fy=shrink, interpolation=cv2.INTER_LINEAR)
import torchvision.transforms as transforms
from torch.autograd import Variable
import scipy.io as sio
from PIL import Image, ImageDraw
from pyramid_vgg import build_vgg
from pyramid import build_sfd
from layers import *
import cv2
import numpy as np
import math
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'
print('Loading model..')
res1_net = build_sfd('test', 640, 2)
#res2_net = build_sfd('test', 640, 2)
vgg1_net = build_vgg('test', 640, 2)
#vgg2_net = build_vgg('test', 640, 2)
net1 = res1_net
#net2 = res2_net
net3 = vgg1_net
#net4 = vgg2_net
net1.load_state_dict(torch.load('./weights/Res50_pyramid_75000.pth'))
#net2.load_state_dict(torch.load('./weights/Res50_pbrush_30000.pth'))
net3.load_state_dict(torch.load('./weights/Vgg16_pyramid_70000.pth'))
#net4.load_state_dict(torch.load('./weights/Res50_vgg2_60000.pth'))
netlist = [net1, net3] # multi-model testing
def main():
vid = cv2.VideoCapture(video_path)
if not vid.isOpened():
raise IOError("couldn't open the video")
test_counter = 0
while True:
return_value, frame = vid.read()
height,width = frame.shape[:2]
print('Loading model: our_ucsd ..')
ssd_net = build_sfd('test', 640, 2)
net = ssd_net
net.load_state_dict(torch.load(args.our))
net.cuda()
net.eval()
print('Finished loading our model!')
image = frame
image = cv2.resize(image, (2*width, 2*height), interpolation=cv2.INTER_CUBIC)
max_im_shrink = (0x7fffffff / 200.0 / (image.shape[0] * image.shape[1])) ** 0.5 # the max size of input image for caffe
max_im_shrink = 3 if max_im_shrink > 3 else max_im_shrink
shrink = max_im_shrink if max_im_shrink < 1 else 1
det0 = detect_face(image, shrink) # origin test
det1 = flip_test(image, shrink) # flip test
[det2, det3] = multi_scale_test(image, max_im_shrink)#min(2,1400/min(image.shape[0],image.shape[1]))) #multi-scale test
det4 = multi_scale_test_pyramid(image, max_im_shrink)
#########################################################################################################################
print('Loading model: general ..')
ssd_net = build_sfd('test', 640, 2)
net = ssd_net
net.load_state_dict(torch.load(args.general))
net.cuda()
net.eval()
print("Finished loading general's model!")
image = cv2.imread(test_path, cv2.IMREAD_COLOR)
height, width = image.shape[:2]
image = cv2.resize(image, (2*width, 2*height), interpolation=cv2.INTER_CUBIC)
max_im_shrink = (0x7fffffff / 200.0 / (image.shape[0] * image.shape[1])) ** 0.5 # the max size of input image for caffe
max_im_shrink = 3 if max_im_shrink > 3 else max_im_shrink
shrink = max_im_shrink if max_im_shrink < 1 else 1
det5 = detect_face(image, shrink) # origin test
det6 = flip_test(image, shrink) # flip test
[det7, det8] = multi_scale_test(image, max_im_shrink)#min(2,1400/min(image.shape[0],image.shape[1]))) #multi-scale test
det9 = multi_scale_test_pyramid(image, max_im_shrink)
det = np.row_stack((det0, det1, det2, det3, det4, det5, det6, det7, det8, det9))
dets = bbox_vote(det)
dets[:, 0] = dets[:, 0] / 2
dets[:, 1] = dets[:, 1] / 2
dets[:, 2] = dets[:, 2] / 2
dets[:, 3] = dets[:, 3] / 2
#write_to_txt(fw, dets, args.probability)
draw_bboxes(dets,args.probability)
test_counter += 1
if test_counter >=16:
break
print("test_counter: " + str(test_counter))
