Spaces:
Runtime error
Runtime error
import os | |
import subprocess | |
import pickle | |
import cv2 | |
from moviepy.editor import VideoFileClip, ImageSequenceClip | |
import shutil | |
import re | |
# 设置环境变量 | |
os.environ['CUDA_VISIBLE_DEVICES'] = '0' | |
def uninstall_tensorflow(): | |
try: | |
result = subprocess.run(['pip', 'uninstall', '-y', 'tensorflow'], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) | |
print("卸载成功") | |
print("输出:", result.stdout.decode('utf-8')) | |
except subprocess.CalledProcessError as e: | |
print("卸载失败") | |
print("错误:", e.stderr.decode('utf-8')) | |
# 调用函数 | |
uninstall_tensorflow() | |
def set_prompt_and_get_coordinates(output_video, texts=['men', 'the table']): | |
if isinstance(texts, str): | |
texts = texts.split(',') | |
texts = [text.strip() for text in texts] | |
print(texts) | |
with open('/workspace/florence-sam/texts.pkl', 'wb') as file: | |
pickle.dump(texts, file) | |
with open('/workspace/florence-sam/output_video.pkl', 'wb') as file: | |
pickle.dump(output_video, file) | |
command = ['python', '/workspace/florence-sam/1.py'] | |
all_ok_bboxes = subprocess.run(command, capture_output=True, text=True) | |
return all_ok_bboxes | |
def run_sam2(output_video): | |
script_path = '/workspace/florence-sam/2.py' | |
command = ['python3', script_path] | |
sam2_output = subprocess.run(command, capture_output=True, text=True) | |
return sam2_output | |
def create_video_with_audio(image_folder, input_video_path): | |
image_files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))] | |
def natural_sort_key(s, _nsre=re.compile('([0-9]+)')): | |
return [int(text) if text.isdigit() else text.lower() for text in re.split(_nsre, s)] | |
image_files.sort(key=natural_sort_key) | |
if image_files: | |
first_image = cv2.imread(os.path.join(image_folder, image_files[0])) | |
height, width, layers = first_image.shape | |
else: | |
raise ValueError("No valid images found in the folder after skipping the first one.") | |
cap = cv2.VideoCapture(input_video_path) | |
fps = cap.get(cv2.CAP_PROP_FPS) | |
cap.release() | |
image_paths = [os.path.join(image_folder, img) for img in image_files] | |
clip = ImageSequenceClip(image_paths, fps=fps) | |
audio_clip = VideoFileClip(input_video_path).audio | |
final_clip = clip.set_audio(audio_clip) | |
output_video_path = os.path.join('/workspace/sam2_videos', os.path.basename(input_video_path)) | |
os.makedirs(os.path.dirname(output_video_path), exist_ok=True) | |
final_clip.write_videofile(output_video_path, codec='libx264') | |
print(f"Video created successfully: {output_video_path}") | |
return output_video_path | |
def process_videos(source_dir, target_dir, texts="men, the table"): | |
os.makedirs(target_dir, exist_ok=True) | |
for video_file in os.listdir(source_dir): | |
if video_file.endswith(('.mp4', '.avi', '.mov')): | |
video_path = os.path.join(source_dir, video_file) | |
print(f"Processing video: {video_path}") | |
result = set_prompt_and_get_coordinates(video_path, texts) | |
print(result) | |
result = run_sam2(video_path) | |
print(result) | |
with open('/workspace/florence-sam/output_video.pkl', 'rb') as file: | |
output_video = pickle.load(file) | |
create_video_with_audio('/workspace/output', output_video) | |
# 移动原视频到目标目录 | |
shutil.move(video_path, os.path.join(target_dir, video_file)) | |
if __name__ == "__main__": | |
source_dir = '/workspace/transvnet/o_videos' | |
target_dir = '/workspace/alreayvideos' | |
texts = "men, the table" | |
process_videos(source_dir, target_dir, texts) |