Spaces:
Running
Running
File size: 2,217 Bytes
5f028d6 |
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 |
import sys
import os
sys.path.insert(0, os.path.dirname(__file__) + '/../..')
import argparse
import numpy as np
from glob import glob
from lib.pipeline.tools import detect_track
from natsort import natsorted
import subprocess
def extract_frames(video_path, output_folder):
if not os.path.exists(output_folder):
os.makedirs(output_folder)
command = [
'ffmpeg',
'-i', video_path,
'-vf', 'fps=30',
'-start_number', '0',
os.path.join(output_folder, '%04d.jpg')
]
subprocess.run(command, check=True)
def detect_track_video(args):
file = args.video_path
root = os.path.dirname(file)
seq = os.path.basename(file).split('.')[0]
seq_folder = f'{root}/{seq}'
img_folder = f'{seq_folder}/extracted_images'
os.makedirs(seq_folder, exist_ok=True)
os.makedirs(img_folder, exist_ok=True)
print(f'Running detect_track on {file} ...')
##### Extract Frames #####
imgfiles = natsorted(glob(f'{img_folder}/*.jpg'))
# print(imgfiles[:10])
if len(imgfiles) > 0:
print("Skip extracting frames")
else:
_ = extract_frames(file, img_folder)
imgfiles = natsorted(glob(f'{img_folder}/*.jpg'))
##### Detection + Track #####
print('Detect and Track ...')
start_idx = 0
end_idx = len(imgfiles)
if os.path.exists(f'{seq_folder}/tracks_{start_idx}_{end_idx}/model_boxes.npy'):
print(f"skip track for {start_idx}_{end_idx}")
return start_idx, end_idx, seq_folder, imgfiles
os.makedirs(f"{seq_folder}/tracks_{start_idx}_{end_idx}", exist_ok=True)
boxes_, tracks_ = detect_track(imgfiles, thresh=0.2)
np.save(f'{seq_folder}/tracks_{start_idx}_{end_idx}/model_boxes.npy', boxes_)
np.save(f'{seq_folder}/tracks_{start_idx}_{end_idx}/model_tracks.npy', tracks_)
return start_idx, end_idx, seq_folder, imgfiles
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--img_focal", type=float)
parser.add_argument("--video_path", type=str, default='')
parser.add_argument("--input_type", type=str, default='file')
args = parser.parse_args()
detect_track_video(args) |