florence-sam-tencent / all_in_one.py
supersolar's picture
Update all_in_one.py
68f1d85 verified
# @title #设置提示词,获取坐标
import subprocess
import subprocess
def uninstall_tensorflow():
try:
# 使用 subprocess.run 执行 pip uninstall 命令
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()
import os
# 设置环境变量
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
import pickle
import cv2
import os
import re
import pickle
from moviepy.editor import VideoFileClip, ImageSequenceClip
import subprocess
def set_prompt_and_get_coordinates(output_video, texts=['men', 'the table']):
if isinstance(texts, str):
texts = texts.split(',') # Assuming prompts are separated by commas
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)
#print(all_ok_bboxes)
return all_ok_bboxes
# 示例调用
output_video = '/workspace/transvnet/o_videos/3.mp4'
texts="men, the table"
result = set_prompt_and_get_coordinates(output_video,texts)
print(result)
# @title #sam2处理
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)
#print(sam2_output)
return sam2_output
# 示例调用
#output_video = 'path/to/output_video.mp4'
result = run_sam2(output_video)
print(result)
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:
image_files = image_files[1:]
'''
# 读取第一张图像以获取尺寸
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('/tmp/gradio/', os.path.basename(input_video_path))
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
# 示例调用
image_folder = '/workspace/output'
with open('/workspace/florence-sam/output_video.pkl', 'rb') as file:
output_video = pickle.load(file)
print(output_video)
input_video_path = output_video
create_video_with_audio(image_folder, input_video_path)
''''''