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))