mirror of
https://github.com/hpcaitech/Open-Sora.git
synced 2026-04-18 16:49:41 +02:00
31 lines
1,011 B
Python
31 lines
1,011 B
Python
# Originally developed by https://github.com/Vchitect/VBench based on https://github.com/facebookresearch/co-tracker.
|
|
|
|
import argparse
|
|
from .camera_motion import compute_camera_motion
|
|
import pandas as pd
|
|
from typing import List
|
|
import ptvsd
|
|
|
|
|
|
def process(paths: List[str], threshold: float) -> List[str]:
|
|
device = "cuda"
|
|
submodules = {'repo': 'facebookresearch/co-tracker', 'model': 'cotracker2'}
|
|
camera_motion_types = compute_camera_motion(device, submodules, paths, factor=threshold)
|
|
return camera_motion_types
|
|
|
|
|
|
def main(args):
|
|
output_file = args.input.replace(".csv", "_cmotion.csv")
|
|
data = pd.read_csv(args.input)
|
|
data["cmotion"] = process(data["path"], args.threshold)
|
|
data.to_csv(output_file, index=False)
|
|
print(f"Output saved to {output_file}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("input", type=str)
|
|
parser.add_argument("--threshold", type=float, default=0.25)
|
|
args = parser.parse_args()
|
|
main(args)
|