def light_test_widerface(): # load net cfg = widerface_640 num_classes = len(WIDERFace_CLASSES) + 1 # +1 background net = build_ssd('test', cfg['min_dim'], num_classes) # initialize SSD net.load_state_dict(torch.load(args.trained_model)) net.cuda() net.eval() print('Finished loading model!') # load data testset = WIDERFaceDetection(args.widerface_root, 'val', None, WIDERFaceAnnotationTransform()) # evaluation cuda = args.cuda transform = TestBaseTransform((104, 117, 123)) thresh = cfg['conf_thresh'] print(thresh) save_path = args.save_folder num_images = len(testset) for i in range(num_images): img = testset.pull_image(i) img_id, annotation = testset.pull_anno(i) event = testset.pull_event(i) print('Testing image {:d}/{:d} {}....'.format(i + 1, num_images, img_id)) shrink = 1 det = infer(net, img, transform, thresh, cuda, shrink) if not os.path.exists(save_path + event): os.makedirs(save_path + event) f = open(save_path + event + '/' + img_id.split(".")[0] + '.txt', 'w') write_to_txt(f, np.array(det), event, img_id)
def light_test_oneimage(): # load net cfg = widerface_640 num_classes = len(WIDERFace_CLASSES) + 1 # +1 background net = build_ssd('test', cfg['min_dim'], num_classes) # initialize SSD #net = nn.DataParallel(net) net.load_state_dict(torch.load(args.trained_model)) net.cuda() net.eval() print('Finished loading model!') # evaluation cuda = args.cuda transform = TestBaseTransform((104, 117, 123)) thresh = cfg['conf_thresh'] # load data #path = './data/worlds-largest-selfie-biggest.jpg' #path = "./data/worlds-largest-selfie.jpg" path = "./data/yuebing.jpg" img_id = 'result' img = cv2.imread(path, cv2.IMREAD_COLOR) shrink = 1 det = infer(net, img, transform, thresh, cuda, shrink) vis_detections(img, det, img_id, 0.6)
def main(): # load net shrink = 1 cfg = widerface_640 num_classes = len(WIDERFace_CLASSES) + 1 # +1 background net = build_ssd('test', cfg['min_dim'], num_classes) # initialize SSD # net = nn.DataParallel(net) net.load_state_dict(torch.load(args.trained_model)) cuda = args.cuda if cuda: net.cuda() else: net.cpu() net.eval() print('Finished loading model!') # evaluation transform = TestBaseTransform((104, 117, 123)) thresh = cfg['conf_thresh'] video_folder_list = sorted([ filename for filename in os.listdir(args.input_folder) if not filename.startswith('.') ]) #隠しフォルダを除く for video in video_folder_list: ext = os.path.splitext(os.path.split(video)[1])[1] print(ext) save_folder = args.input_folder.replace("input", "output") if not os.path.exists(save_folder): os.makedirs(save_folder) if ext == ".mp4" or ext == ".MP4": save_folder_tmp = save_folder + "tmp/" save_folder_video = save_folder + "videos/" if not os.path.exists(save_folder_tmp): os.makedirs(save_folder_tmp) os.makedirs(save_folder_video) frame_id = 0 video_name = os.path.splitext(os.path.split(video)[1])[0] # GoProなどをマスクする際には音声情報によりcv2.VideoCaptureが途中で終了することがある. よってffmpegを用いて音声を削除するコードを挟んでいる. if args.delete_audio: import subprocess command_delete_audio = [ "ffmpeg", "-y", "-i", args.input_folder + video, "-vcodec", "copy", "-map", "0:0", args.input_folder + video.replace(".", "_an.") ] subprocess.run(command_delete_audio) video = video.replace(".", "_an.") cap = cv2.VideoCapture(args.input_folder + video) original_fps = cap.get(cv2.CAP_PROP_FPS) while True: _, frame = cap.read() print(_) if not _: break frame_id += 1 # if frame_id<=1358: # continue det = infer(net, frame, transform, thresh, cuda, shrink) cv2.imwrite( save_folder_tmp + video_name + '_' + '{:09d}.jpg'.format(frame_id), frame) print('prossing:', frame_id) if det[0][0] == 0.1: cv2.imwrite( save_folder_tmp + video_name + '_' + '{:09d}.jpg'.format(frame_id), frame) save_mosaiced_img(frame, det, save_folder_tmp, video_name, frame_id, scale=args.mosaic_scale, thresh=args.visual_threshold, kernel_size=args.mosaic_density) save_folder_list = sorted([ filename for filename in os.listdir(save_folder_tmp) if not filename.startswith('.') ]) #隠しフォルダを除く print(len(save_folder_list)) print(save_folder_video) make_video_from_images(save_folder_tmp, save_folder_list, os.path.join(save_folder_video, video_name + '.mp4'), fps=original_fps) delete_imgs(save_folder_tmp) elif ext == ".jpg" or ext == ".png": save_folder_img = save_folder + "imgs/" if not os.path.exists(save_folder_img): os.makedirs(save_folder_img) video_name = os.path.splitext(os.path.split(video)[1])[0] frame_id = -1 img_path = args.input_folder + video frame = cv2.imread(img_path) det = infer(net, frame, transform, thresh, cuda, shrink) if det[0][0] == 0.1: cv2.imwrite(save_folder_img + video_name + '.jpg', frame) save_mosaiced_img(frame, det, save_folder_img, video_name, frame_id, scale=args.mosaic_scale, thresh=args.visual_threshold, kernel_size=args.mosaic_density) delete_imgs(save_folder_tmp)
cap.get(cv2.CAP_PROP_FRAME_WIDTH) // args.reduce_scale), int( cap.get(cv2.CAP_PROP_FRAME_HEIGHT) // args.reduce_scale)) out = cv2.VideoWriter(args.out, fourcc, cap.get(cv2.CAP_PROP_FPS), out_size) total_frame = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) f_count = 0 e_time = None s_time = None # load net cfg = widerface_640 num_classes = len(WIDERFace_CLASSES) + 1 # +1 background net = build_ssd('test', cfg['min_dim'], num_classes) # initialize SSD # net = nn.DataParallel(net) net.load_state_dict(torch.load(args.trained_model)) net.cuda() net.eval() print('Finished loading model!') # evaluation cuda = args.cuda transform = TestBaseTransform((104, 117, 123)) thresh = cfg['conf_thresh'] shrink = 1 while cap.isOpened(): if args.verbose > 0 and e_time is not None: