66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
from lib.qwen_fun import hazard_inspection, search_knowledge_base
|
||
from encodings.punycode import T
|
||
"""
|
||
测试 在给定标注框与类别，原始视频经过转换之后，AI能否准确识别物体特征
|
||
"""
|
||
|
||
from tkinter import N
|
||
from datetime import datetime
|
||
import json
|
||
import os
|
||
from lib.qwen_fun import chat, get_annnotated_frame_for_ai, get_unique_track_id_count, save_json_to_file, search_knowledge_base, upload_files_and_get_urls_concurrently
|
||
from lib.qwen_fun_vid import create_mian_vid_for_ai, frame_all_to_obj_vid
|
||
from lib.sam3 import SAM3
|
||
from ultralytics.models.sam import SAM3VideoSemanticPredictor
|
||
import cv2
|
||
import json
|
||
import numpy as np
|
||
from pathlib import Path
|
||
|
||
def run(output_dir, vid_path, interval: int = 5):
|
||
if output_dir:
|
||
os.makedirs(output_dir, exist_ok=True) # exist_ok=True 防止重复创建报错
|
||
|
||
annotated_frames: SAM3 = SAM3()
|
||
annotated_frames.run(vid_path, output_dir, "lib/class_list/xiaofangz.json")
|
||
# annotated_frames.load_from_json(f"{output_dir}/frame_all.json")
|
||
# print(annotated_frames.data)
|
||
|
||
# 提取ai能看到的部分
|
||
ai_frames: dict = get_annnotated_frame_for_ai(annotated_frames.data(), interval)
|
||
save_json_to_file(ai_frames, f"{output_dir}/frame_all_ai.json")
|
||
|
||
|
||
if vid_path.startswith("oss"):
|
||
video_url = vid_path
|
||
else:
|
||
video_url: str|None = upload_files_and_get_urls_concurrently(
|
||
file_path_list=[vid_path],
|
||
max_workers=8
|
||
)[0]
|
||
if video_url:
|
||
with open(f"{output_dir}/video_url.json", "w") as f:
|
||
f.write(video_url)
|
||
|
||
if video_url is None:
|
||
raise ValueError("视频上传失败，无法获取 URL")
|
||
|
||
result = hazard_inspection(ai_frames, video_url)[1]
|
||
|
||
with open(f"{output_dir}/hazard_inspection.json", "w", encoding="utf-8") as f:
|
||
f.write(result)
|
||
|
||
# 启动应用
|
||
if __name__ == "__main__":
|
||
VID_NAME: str = "Peihdianhxiang"
|
||
VID_END: str = ".mp4"
|
||
interval: int = int(30 / 5)
|
||
|
||
output_dir: str = f"output/{VID_NAME}"
|
||
vid_path: str = f"./input/{VID_NAME}{VID_END}"
|
||
|
||
run(output_dir, vid_path, interval)
|
||
|
||
|
||
|