Speed up HF analysis progress reporting

RAG/Knowledge_Database/RAGFunc.py

@@ -168,6 +168,7 @@ def get_video_ori_keywords(
     language='zh',
     show=False,
     template_keyframes_dir = None,
+    progress_callback=None,
 ) -> dict:
     """
     Multi-stage multimodal assessment for archery posture.
@@ -191,6 +192,13 @@ def get_video_ori_keywords(
             return "en"
         raise ValueError("language must be 'en' or 'zh'")
 
+    def _notify_progress(stage: str, message: str) -> None:
+        if callable(progress_callback):
+            try:
+                progress_callback(stage, message)
+            except Exception:
+                pass
+
     def _chat_completion_output_text(resp) -> str:
         if resp is None or not getattr(resp, "choices", None):
             raise ValueError("Chat completion returned empty response.")
@@ -539,13 +547,22 @@ def get_video_ori_keywords(
         image_width=1920,
         image_height=1080,
         draw_math_feature_points=show,
+        progress_callback=progress_callback,
+    )
+    _notify_progress("selecting_keyframes", "正在筛选关键帧")
+    base_keyframes = extract_keyframes_with_ruptures_poseparts_2d(
+        normalized_data,
+        k=target_k + 3,
+        print_all_frame_scores=show,
+        print_selection_debug=show,
     )
-    base_keyframes = extract_keyframes_with_ruptures_poseparts_2d(normalized_data, k=target_k + 3)
+    _notify_progress("refining_keyframes", "正在细化关键帧")
     key_frame_lists = refine_keyframes_with_absdiff(
         video_path=video_path,
         keyframe_result=base_keyframes,
         k=target_k,
     )
+    _notify_progress("rendering_keyframes", "正在整理关键帧输入")
     keyframes_dict = extract_show_keyframes_by_index(video_path, key_frame_lists, show=show)
     keyframe_image_items = list(keyframes_dict.get("openai_input_images", []) or [])
     if not keyframe_image_items:
@@ -593,6 +610,7 @@ def get_video_ori_keywords(
         # + f"- min_x_diff_avg: {metrics['min_x_diff_avg']}"
     )
     if pipeline == 3:
+        _notify_progress("generating_assessment", "正在生成动作评估")
         p3_system_prompt = system_prompt + system_metrics_rubric_map[lang]
         user_text = ordering_note #+ metrics_text
         raw_user_content = [{"type": "input_text", "text": user_text}] + keyframe_image_items
@@ -607,6 +625,7 @@ def get_video_ori_keywords(
         return _parse_assessment_payload(_chat_completion_output_text(resp))
 
     if pipeline == 4:
+        _notify_progress("loading_template_keyframes", "正在加载模板关键帧")
         repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
         default_template_dir = os.path.join(repo_root, "output_keyframes")
         template_dir = template_keyframes_dir if template_keyframes_dir else default_template_dir
@@ -629,6 +648,7 @@ def get_video_ori_keywords(
             + template_keyframe_items
         )
         chat_user_content = _to_chat_content(raw_user_content)
+        _notify_progress("generating_assessment", "正在生成动作评估")
         resp = client.chat.completions.create(
             model=model_name,
             messages=[

RAG/tokenize_search.py

@@ -2,7 +2,7 @@ from RAG.Knowledge_Database.RAGFunc import *
 from RAG.Knowledge_Database.AIdbconfig import session, session_en
 from RAG.Knowledge_Database.AI_dbmanager import KnowledgeDB
 
-def Tokenize_SearchKeyword(video_path, pipeline=1, subpipeline=3, language='zh',show=False):
+def Tokenize_SearchKeyword(video_path, pipeline=1, subpipeline=3, language='zh', show=False, progress_callback=None):
     """
     Extract keywords from video and search knowledge database
     
@@ -15,7 +15,13 @@ def Tokenize_SearchKeyword(video_path, pipeline=1, subpipeline=3, language='zh',
         List of keywords from knowledge database
     """
     if pipeline == 1:
-        answer_content = get_video_ori_keywords(video_path, pipeline=subpipeline, language=language,show=show)
+        answer_content = get_video_ori_keywords(
+            video_path,
+            pipeline=subpipeline,
+            language=language,
+            show=show,
+            progress_callback=progress_callback,
+        )
         total_score = answer_content["total_score"]
         head_score  = answer_content["head_score"]
         hand_score  = answer_content["hand_score"]
@@ -43,6 +49,11 @@ def Tokenize_SearchKeyword(video_path, pipeline=1, subpipeline=3, language='zh',
         db = KnowledgeDB(session=session_en if language == 'en' else session)
         top_k = 34 if language == 'en' else 34
         #TODO 看embedding模型能不能换成qwen-embedding?
+        if callable(progress_callback):
+            try:
+                progress_callback("retrieving_knowledge", "正在检索技术知识库")
+            except Exception:
+                pass
         return score_dict, comment, db.from_video_search(query_vec=query_embeddings[0], model_name='ali-text-embedding-v3', top_k=top_k)
     
     elif pipeline == 2:
@@ -51,4 +62,4 @@ def Tokenize_SearchKeyword(video_path, pipeline=1, subpipeline=3, language='zh',
         if language == 'en':
             return db.from_video_search(query_vec=video_token, model_name='languagebind', top_k=34)
         else:
-            return db.from_video_search(query_vec=video_token, model_name='languagebind', top_k=17)
+            return db.from_video_search(query_vec=video_token, model_name='languagebind', top_k=17)

RTMPose/Bone_Feature_Extract.py

@@ -5,9 +5,13 @@ import ruptures as rpt
 import matplotlib.pyplot as plt
 import os
 import base64
+import threading
 
 from Tools.Exe_dataset.model_config import model_configs
 
+_POSE_TRACKER_CACHE = {}
+_POSE_TRACKER_LOCK = threading.Lock()
+
 COCO133_KPT_IDX = {
     # Body 17
     "left_shoulder": 5,
@@ -71,6 +75,7 @@ def Keypoint_Extract(
     track_conf_kpt_thr=0.2,
     draw_tracked_only=True,
     draw_math_feature_points=False,
+    progress_callback=None,
 ):
     from RTMPose.rtmlib import PoseTracker, Wholebody3d, draw_skeleton
     """
@@ -209,19 +214,39 @@ def Keypoint_Extract(
             )
         return frame_bgr
 
+    def _notify_progress(stage, message):
+        if callable(progress_callback):
+            try:
+                progress_callback(stage, message)
+            except Exception:
+                pass
+
+    def _get_cached_pose_tracker(backend_name, device_name):
+        cache_key = (backend_name, device_name)
+        with _POSE_TRACKER_LOCK:
+            tracker = _POSE_TRACKER_CACHE.get(cache_key)
+            if tracker is None:
+                tracker = PoseTracker(
+                    Wholebody3d,
+                    det_frequency=7,
+                    tracking=False,
+                    backend=backend_name,
+                    device=device_name,
+                )
+                _POSE_TRACKER_CACHE[cache_key] = tracker
+        return tracker
+
     # Keep this False by default; turn on only when diagnosing index mismatches.
     _maybe_print_coco133_index_self_check(enabled=False)
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     backend = 'onnxruntime'
     cap = cv2.VideoCapture(path_to_video)
+    device_name = 'cuda' if device.type == 'cuda' else 'cpu'
 
-    wholebody3d = PoseTracker(
-        Wholebody3d,
-        det_frequency=7,
-        tracking=False,
-        backend=backend,
-        device='cuda' if device.type == 'cuda' else 'cpu')
+    _notify_progress("loading_pose_model", "正在加载姿态模型")
+    wholebody3d = _get_cached_pose_tracker(backend, device_name)
+    _notify_progress("extracting_keypoints", "正在提取人体关键点")
 
     frame_idx = -1
     whole_skeleton_data = []

webapp/backend/app.py

@@ -67,16 +67,26 @@ def _score_payload(scores: dict[str, Any]) -> dict[str, Any]:
     }
 
 
-def analyze_video(video_path: str, language: str, subpipeline: int) -> dict[str, Any]:
+def analyze_video(
+    video_path: str,
+    language: str,
+    subpipeline: int,
+    progress_callback=None,
+) -> dict[str, Any]:
     _ensure_api_key()
+    if callable(progress_callback):
+        progress_callback("starting", "正在准备分析任务")
     scores, raw_keywords, retrieved_result = Tokenize_SearchKeyword(
         video_path=video_path,
         pipeline=1,
         subpipeline=subpipeline,
         language=language,
         show=False,
+        progress_callback=progress_callback,
     )
     keywords = _coerce_keywords(raw_keywords)
+    if callable(progress_callback):
+        progress_callback("writing_assessment", "正在生成评估结论")
     assessment_text = get_response(
         keywords=keywords,
         score_dict=scores,
@@ -116,6 +126,8 @@ class JobRecord:
     filename: str
     language: str
     status: str = "queued"
+    stage: str = "queued"
+    status_message: str = "任务已创建，等待处理"
     created_at: float = field(default_factory=time.time)
     updated_at: float = field(default_factory=time.time)
     error: str | None = None
@@ -152,12 +164,24 @@ class JobStore:
         with self._lock:
             job = self._jobs[job_id]
             job.status = "running"
+            job.stage = "starting"
+            job.status_message = "任务开始执行"
+            job.updated_at = time.time()
+
+    def set_progress(self, job_id: str, stage: str, status_message: str) -> None:
+        with self._lock:
+            job = self._jobs[job_id]
+            job.status = "running"
+            job.stage = stage
+            job.status_message = status_message
             job.updated_at = time.time()
 
     def set_completed(self, job_id: str, result: dict[str, Any]) -> None:
         with self._lock:
             job = self._jobs[job_id]
             job.status = "completed"
+            job.stage = "completed"
+            job.status_message = "分析完成，可以继续提问"
             job.result = result
             job.updated_at = time.time()
 
@@ -165,6 +189,8 @@ class JobStore:
         with self._lock:
             job = self._jobs[job_id]
             job.status = "failed"
+            job.stage = "failed"
+            job.status_message = error
             job.error = error
             job.updated_at = time.time()
 
@@ -194,6 +220,8 @@ def _public_job_payload(record: JobRecord) -> dict[str, Any]:
         "filename": record.filename,
         "language": record.language,
         "status": record.status,
+        "stage": record.stage,
+        "status_message": record.status_message,
         "created_at": record.created_at,
         "updated_at": record.updated_at,
         "error": record.error,
@@ -213,7 +241,15 @@ def _public_job_payload(record: JobRecord) -> dict[str, Any]:
 def _run_job(job_id: str, temp_video_path: str, language: str, subpipeline: int) -> None:
     jobs.set_running(job_id)
     try:
-        result = analyze_video(temp_video_path, language=language, subpipeline=subpipeline)
+        def progress_callback(stage: str, message: str) -> None:
+            jobs.set_progress(job_id, stage, message)
+
+        result = analyze_video(
+            temp_video_path,
+            language=language,
+            subpipeline=subpipeline,
+            progress_callback=progress_callback,
+        )
         jobs.set_completed(job_id, result)
     except Exception as exc:
         jobs.set_failed(job_id, str(exc))

webapp/static/app.js

@@ -29,9 +29,9 @@ function resetResults() {
   activeJobId = null;
   jobIdEl.textContent = "";
   scoresEl.innerHTML = "";
-  keywordsEl.textContent = "分析完成后显示。";
+  keywordsEl.textContent = "分析完成后显示";
   keywordsEl.className = "chips empty";
-  assessmentEl.textContent = "分析完成后显示。";
+  assessmentEl.textContent = "分析完成后显示";
   assessmentEl.className = "answer-box empty";
   chatLog.innerHTML = '<div class="message message-system">分析完成后，可以继续和 SEMA 交互。</div>';
   questionInput.disabled = true;
@@ -57,7 +57,7 @@ function renderScores(scores) {
 
 function renderKeywords(keywords) {
   if (!keywords || keywords.length === 0) {
-    keywordsEl.textContent = "无关键词。";
+    keywordsEl.textContent = "无关键词";
     keywordsEl.className = "chips empty";
     return;
   }
@@ -96,7 +96,10 @@ async function pollJob(jobId) {
     }
 
     if (job.status === "queued" || job.status === "running") {
-      setStatus("running", job.status === "queued" ? "任务已创建，等待处理。" : "正在分析视频，请稍候。");
+      const fallbackText = job.status === "queued"
+        ? "任务已创建，等待处理。"
+        : "正在分析视频，请稍候。";
+      setStatus("running", job.status_message || fallbackText);
       pollHandle = setTimeout(() => pollJob(jobId), 2500);
       return;
     }
@@ -139,7 +142,7 @@ analyzeForm.addEventListener("submit", async (event) => {
     if (!response.ok) {
       throw new Error(payload.detail || "任务创建失败。");
     }
-    setStatus("running", "上传完成，任务已进入队列。");
+    setStatus("running", payload.status_message || "上传完成，任务已进入队列。");
     pollJob(payload.job_id);
   } catch (error) {
     setStatus("error", error.message || "上传失败。");