136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
import os
|
||
|
||
from lib.json_fun import f_detections_to_objects
|
||
from lib.qwen_fun import chat, hazard_inspection, load_json_data, search_knowledge_base, upload_files_and_get_urls_concurrently
|
||
from encodings.punycode import T
|
||
"""
|
||
测试 在给定标注框与类别，原始视频经过转换之后，AI能否准确识别物体特征
|
||
"""
|
||
|
||
from tkinter import N
|
||
from datetime import datetime
|
||
import json
|
||
from lib.qwen_fun_vid import generate_video_to_objects
|
||
from lib.sam3 import SAM3
|
||
from ultralytics.models.sam import SAM3VideoSemanticPredictor
|
||
import json
|
||
import numpy as np
|
||
from pathlib import Path
|
||
|
||
def run():
|
||
vid_name = "santai5" # 视频名称（不含后缀）
|
||
vid_end = ".mp4" # 视频后缀
|
||
use_url_cache = True # 是否使用 URL 缓存，避免重复上传视频
|
||
enable_thinking = False # 是否启用思考模式
|
||
|
||
run_vid_process = False # 是否运行视频处理流程（提取物体视频）
|
||
|
||
input_video_path = f"input/{vid_name}{vid_end}"
|
||
output_dir = f"output/{vid_name}"
|
||
frame_detections_path = f"{output_dir}/frame_detections.json"
|
||
objects_json_path = f"{output_dir}/objects.json"
|
||
vid_dir = f"{output_dir}/obj_vids"
|
||
obj = json.load(open(objects_json_path, "r", encoding="utf-8"))
|
||
class_list = obj["class_list"]
|
||
obj_dict = obj["track_id_list"]
|
||
rule_dict: dict = load_json_data('知识库/rule.json')
|
||
video_url_file = f"{output_dir}/video_url.json"
|
||
|
||
# 提取物体信息
|
||
f_detections_to_objects(
|
||
frame_detections_path,
|
||
objects_json_path
|
||
)
|
||
|
||
if run_vid_process:
|
||
# 生成物体视频
|
||
generate_video_to_objects(
|
||
obj_dict,
|
||
input_video_path,
|
||
output_dir=vid_dir,
|
||
)
|
||
|
||
# 上传视频并获取 URL
|
||
vid_dict = get_vid_dict(vid_dir, obj_dict, video_url_file, use_url_cache)
|
||
|
||
hazard_inspection(
|
||
output_dir,
|
||
obj_dict,
|
||
rule_dict,
|
||
class_list,
|
||
vid_dict,
|
||
fps=30,
|
||
enable_thinking=enable_thinking
|
||
)
|
||
|
||
def get_vid_dict(vid_dir: str, obj_dict: dict, video_url_file: str, use_url_cache: bool) -> dict:
|
||
"""获取视频字典
|
||
参数:
|
||
vid_dir: 物体视频存放目录
|
||
obj_dict: 物品字典，键为物品TrackID，值为物品信息
|
||
video_url_file: 视频URL文件路径
|
||
use_url_cache: 是否使用缓存的URL
|
||
返回:
|
||
视频字典，键为物品TrackID，值为视频本地地址和视频URL
|
||
"""
|
||
|
||
vid_dict = {}
|
||
|
||
# 检查URL是否已存在
|
||
if use_url_cache:
|
||
if os.path.exists(video_url_file):
|
||
try:
|
||
with open(video_url_file, "r", encoding="utf-8") as f:
|
||
vid_dict = json.load(f)
|
||
except Exception as e:
|
||
print(f"读取URL文件失败: {e}")
|
||
for track_id in obj_dict.keys():
|
||
if track_id in vid_dict:
|
||
continue
|
||
vid_path = f"{vid_dir}/obj_{int(track_id):03d}.mp4"
|
||
vid_dict[track_id] = {"vid_path": vid_path, "vid_url": None}
|
||
|
||
# 找到未上传的视频
|
||
unuploaded_vids = []
|
||
for track_id, vid_info in vid_dict.items():
|
||
if vid_info["vid_url"] is None:
|
||
unuploaded_vids.append(vid_info["vid_path"])
|
||
|
||
# 上传未上传的视频并获取 URL
|
||
uploaded_urls = upload_files_and_get_urls_concurrently(
|
||
file_path_list=unuploaded_vids,
|
||
max_workers=8
|
||
)
|
||
|
||
# 更新视频URL
|
||
for track_id, vid_info in vid_dict.items():
|
||
if vid_info["vid_url"] is None:
|
||
idx = unuploaded_vids.index(vid_info["vid_path"])
|
||
vid_info["vid_url"] = uploaded_urls[idx]
|
||
|
||
# 保存字典为 JSON 文件
|
||
with open(video_url_file, "w", encoding="utf-8") as f:
|
||
json.dump(vid_dict, f, ensure_ascii=False, indent=4)
|
||
|
||
return vid_dict
|
||
|
||
|
||
# 启动应用
|
||
if __name__ == "__main__":
|
||
# objects_json_path = "output/santai5/objects.json"
|
||
# input_video_path = "input/santai5.mp4"
|
||
# obj_dict = json.load(open(objects_json_path, "r", encoding="utf-8"))["track_id_list"]
|
||
|
||
# # 记录开始时间
|
||
# start_time = datetime.now()
|
||
# generate_video_to_objects(
|
||
# obj_dict,
|
||
# input_video_path,
|
||
# output_dir="output/santai5/obj_vids",
|
||
# )
|
||
|
||
run()
|
||
|
||
|
||
|