| 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 |
| import spaces |
|
|
|
|
| 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) |
|
|
| @spaces.GPU(duration=80) |
| 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) |
