|
|
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} ...') |
|
|
|
|
|
|
|
|
imgfiles = natsorted(glob(f'{img_folder}/*.jpg')) |
|
|
|
|
|
if len(imgfiles) > 0: |
|
|
print("Skip extracting frames") |
|
|
else: |
|
|
_ = extract_frames(file, img_folder) |
|
|
imgfiles = natsorted(glob(f'{img_folder}/*.jpg')) |
|
|
|
|
|
|
|
|
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) |
