PromptTrack / app.py
SophieNgo's picture
Update app.py
2cf6d1b verified
import spaces
import gradio as gr
import shutil
import os
import subprocess
import sys
# Run the .bat file before launching the app
try:
import PromptTrack
except ImportError:
print("PromptTrack not found. Installing...")
subprocess.run([sys.executable, "-m", "pip", "install",
"--index-url", "https://test.pypi.org/simple/",
"--extra-index-url", "https://pypi.org/simple/",
"PromptTrack"], check=True)
subprocess.run([sys.executable, "-m", "pip", "install",
"--no-deps", "bytetracker"], check=True)
import PromptTrack # Retry import after installation
from PromptTrack import PromptTracker
tracker = PromptTracker()
#@spaces.GPU(duration=300)
def process_video(video_path, prompt):
import torch
print(f"Is CUDA available: {torch.cuda.is_available()}")
# True
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
# Tesla T4
detection_threshold=0.3
track_thresh=0.4
match_thresh=1
max_time_lost=float("inf")
nbr_frames_fixing=800
output_video = video_path.split('mp4')[0]+"_with_id.mp4" # Placeholder for processed video
output_file = video_path.split('mp4')[0]+"_mot_.json" # Tracking result
output_file_2 = video_path.split('mp4')[0]+"_object_detection.json" # detection results
video_file = video_path
tracker.detect_objects(video_file, prompt=prompt, nms_threshold=0.8, detection_threshold=detection_threshold, detector="OWL-VITV2")
tracker.process_mot(video_file, fixed_parc=True, track_thresh=track_thresh, match_thresh=match_thresh, frame_rate=25, max_time_lost=max_time_lost, nbr_frames_fixing=nbr_frames_fixing)
tracker.read_video_with_mot(video_file, fps=25)
"""output_video = "output.mp4" # Placeholder for processed video
output_file = "output.txt" # Placeholder for generated file
# Copy the input video to simulate processing
shutil.copy(video_path.name, output_video)
# Create an output text file with the prompt content
with open(output_file, "w") as f:
f.write(f"User Prompt: {prompt}\n")
"""
return output_video, output_file
# Define Gradio interface
iface = gr.Interface(
fn=process_video,
inputs=[gr.File(label="Upload Video"), gr.Textbox(placeholder="Enter your prompt")],
outputs=[gr.Video(), gr.File(label="Generated File")],
title="Video Processing App",
description="Upload a video and enter a prompt. The app will return the processed video and a generated file."
)
# Launch the app
if __name__ == "__main__":
iface.launch(share=True)