Spaces:
Running
Running
File size: 4,280 Bytes
d015578 |
1 2 3 4 5 6 7 8 9 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import os
import cv2
import pkg_resources
# My libs
import spiga.demo.analyze.track.get_tracker as tr
import spiga.demo.analyze.extract.spiga_processor as pr_spiga
from spiga.demo.analyze.analyzer import VideoAnalyzer
from spiga.demo.visualize.viewer import Viewer
# Paths
video_out_path_dft = pkg_resources.resource_filename('spiga', 'demo/outputs')
if not os.path.exists(video_out_path_dft):
os.makedirs(video_out_path_dft)
def main():
import argparse
pars = argparse.ArgumentParser(description='Face App')
pars.add_argument('-i', '--input', type=str, default='0', help='Video input')
pars.add_argument('-d', '--dataset', type=str, default='wflw',
choices=['wflw', '300wpublic', '300wprivate', 'merlrav'],
help='SPIGA pretrained weights per dataset')
pars.add_argument('-t', '--tracker', type=str, default='RetinaSort',
choices=['RetinaSort', 'RetinaSort_Res50'], help='Tracker name')
pars.add_argument('-sh', '--show', nargs='+', type=str, default=['fps', 'face_id', 'landmarks', 'headpose'],
choices=['fps', 'bbox', 'face_id', 'landmarks', 'headpose'],
help='Select the attributes of the face to be displayed ')
pars.add_argument('-s', '--save', action='store_true', help='Save record')
pars.add_argument('-nv', '--noview', action='store_false', help='Do not visualize the window')
pars.add_argument('--outpath', type=str, default=video_out_path_dft, help='Video output directory')
pars.add_argument('--fps', type=int, default=30, help='Frames per second')
pars.add_argument('--shape', nargs='+', type=int, help='Visualizer shape (W,H)')
args = pars.parse_args()
if args.shape:
if len(args.shape) != 2:
raise ValueError('--shape requires two values: width and height. Ej: --shape 256 256')
else:
video_shape = tuple(args.shape)
else:
video_shape = None
if not args.noview and not args.save:
raise ValueError('No results will be saved neither shown')
video_app(args.input, spiga_dataset=args.dataset, tracker=args.tracker, fps=args.fps,
save=args.save, output_path=args.outpath, video_shape=video_shape, visualize=args.noview, plot=args.show)
def video_app(input_name, spiga_dataset=None, tracker=None, fps=30, save=False,
output_path=video_out_path_dft, video_shape=None, visualize=True, plot=()):
# Load video
try:
capture = cv2.VideoCapture(int(input_name))
video_name = None
if not visualize:
print('WARNING: Webcam must be visualized in order to close the app')
visualize = True
except:
try:
capture = cv2.VideoCapture(input_name)
video_name = input_name.split('/')[-1][:-4]
except:
raise ValueError('Input video path %s not valid' % input_name)
if capture is not None:
# Initialize viewer
if video_shape is not None:
vid_w, vid_h = video_shape
else:
vid_w, vid_h = capture.get(3), capture.get(4)
viewer = Viewer('face_app', width=vid_w, height=vid_h, fps=fps)
if save:
viewer.record_video(output_path, video_name)
# Initialize face tracker
faces_tracker = tr.get_tracker(tracker)
faces_tracker.detector.set_input_shape(capture.get(4), capture.get(3))
# Initialize processors
processor = pr_spiga.SPIGAProcessor(dataset=spiga_dataset)
# Initialize Analyzer
faces_analyzer = VideoAnalyzer(faces_tracker, processor=processor)
# Convert FPS to the amount of milliseconds that each frame will be displayed
if visualize:
viewer.start_view()
while capture.isOpened():
ret, frame = capture.read()
if ret:
# Process frame
faces_analyzer.process_frame(frame)
# Show results
key = viewer.process_image(frame, drawers=[faces_analyzer], show_attributes=plot)
if key:
break
else:
break
capture.release()
viewer.close()
if __name__ == '__main__':
main()
|