fix history video display

app_utils.py

@@ -13,6 +13,7 @@ from collections import defaultdict
 import shutil
 from urllib.parse import urljoin
 import oss2
+from natsort import natsorted
 
 TMP_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
 os.makedirs(TMP_ROOT, exist_ok=True)
@@ -168,7 +169,7 @@ OSS_CONFIG = {
 }
 
 auth = oss2.Auth(OSS_CONFIG["access_key_id"], OSS_CONFIG["access_key_secret"])
-bucket = oss2.Bucket(auth, OSS_CONFIG["endpoint"], OSS_CONFIG["bucket_name"])
+bucket = oss2.Bucket(auth, OSS_CONFIG["endpoint"], OSS_CONFIG["bucket_name"], enable_crc=False)
 
 
 def list_oss_files(folder_path: str) -> List[str]:
@@ -179,9 +180,15 @@ def list_oss_files(folder_path: str) -> List[str]:
             files.append(obj.key)
     return sorted(files, key=lambda x: os.path.splitext(x)[0]) 
 
-def download_oss_file(oss_path: str, local_path: str):
-    """从OSS下载文件到本地"""
-    bucket.get_object_to_file(oss_path, local_path)
+def download_oss_file(oss_path: str, local_path: str) -> bool:
+    """从OSS下载文件到本地，返回是否成功"""
+    try:
+        result = bucket.get_object_to_file(oss_path, local_path)
+        print(f"下载: {oss_path}, {result.status}")
+        return result.status == 200
+    except Exception as e:
+        print(f"下载失败: {e}")
+        return False
 
 def oss_file_exists(oss_path):
     try:
@@ -190,6 +197,7 @@ def oss_file_exists(oss_path):
     except Exception as e:
         print(f"Error checking if file exists in OSS: {str(e)}")
         return False
+
 def stream_simulation_results(result_folder: str, task_id: str, request: gr.Request, fps: int = 30):
     """
     流式输出仿真结果，从OSS读取图片
@@ -208,11 +216,11 @@ def stream_simulation_results(result_folder: str, task_id: str, request: gr.Requ
     image_folder = os.path.join(result_folder, "image")
     os.makedirs(image_folder, exist_ok=True)
     frame_buffer: List[np.ndarray] = []
-    frames_per_segment = fps * 2  # 每2秒60帧
+    min_frames_per_segment = fps * 1  # 至少30帧才输出
     processed_files = set()
     width, height = 0, 0
     last_status_check = 0
-    status_check_interval = 1  # 每5秒检查一次后端状态
+    status_check_interval = 5  # 每5秒检查一次后端状态
     max_time = 240
 
     # 创建临时目录存储下载的图片
@@ -245,7 +253,6 @@ def stream_simulation_results(result_folder: str, task_id: str, request: gr.Requ
         try:
             oss_files = list_oss_files(image_folder)
             new_files = [f for f in oss_files if f not in processed_files]
-            has_new_frames = False
             
             for oss_path in new_files:
                 try:
@@ -262,15 +269,13 @@ def stream_simulation_results(result_folder: str, task_id: str, request: gr.Requ
                         
                         frame_buffer.append(frame)
                         processed_files.add(oss_path)
-                        has_new_frames = True
                 except Exception as e:
                     print(f"Error processing {oss_path}: {e}")
 
-            # 如果有新帧且积累够60帧，输出视频片段
-            if has_new_frames and len(frame_buffer) >= frames_per_segment:
-                segment_frames = frame_buffer[:frames_per_segment]
-                frame_buffer = frame_buffer[frames_per_segment:]
-                yield create_video_segment(segment_frames, fps, width, height, request)
+            # 如果有新帧且积累够60帧以上，输出所有当前帧
+            if len(frame_buffer) >= min_frames_per_segment:
+                yield create_video_segment(frame_buffer, fps, width, height, request)
+                frame_buffer = []  # 清空缓冲区
 
         except Exception as e:
             print(f"Error accessing OSS: {e}")
@@ -279,7 +284,6 @@ def stream_simulation_results(result_folder: str, task_id: str, request: gr.Requ
     
     if max_time <= 0:
         raise gr.Error("timeout 240s")
-
 def create_video_segment(frames: List[np.ndarray], fps: int, width: int, height: int, req: gr.Request) -> str:
     """创建视频片段"""
     user_dir = os.path.join(TMP_ROOT, str(req.session_hash))
@@ -424,6 +428,76 @@ def convert_to_h264(video_path):
     except Exception as e:
         raise gr.Error(f"转换过程中发生错误: {str(e)}")
 
+
+def generate_whole_video(task_id: str, request: gr.Request, fps: int = 30) -> str:
+    """
+    从图片序列生成完整视频
+    
+    Args:
+        task_id: 任务ID
+        fps: 视频帧率，默认为30
+    
+    Returns:
+        生成的视频文件路径
+    """
+    frame_buffer: List[np.ndarray] = []
+    user_dir = os.path.join(TMP_ROOT, str(request.session_hash))
+    image_folder = os.path.join(user_dir, task_id, "tasks", "images")
+    
+
+    # 确保输出目录存在
+    result_folder = os.path.join(user_dir, task_id, "tasks", "video")
+    os.makedirs(result_folder, exist_ok=True)
+    
+    # 获取所有图片文件并按自然顺序排序
+    image_files = [f for f in os.listdir(image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+    image_files = natsorted(image_files)  # 自然排序处理类似 'frame1, frame2, ..., frame10' 的情况
+    
+    if not image_files:
+        raise ValueError(f"No image files found in {image_folder}")
+    
+    # 初始化视频尺寸
+    width, height = 0, 0
+    
+    for img_file in image_files:
+        img_path = os.path.join(image_folder, img_file)
+        try:
+            frame = cv2.imread(img_path)
+            if frame is not None:
+                if width == 0:  # 第一次获取图像尺寸
+                    height, width = frame.shape[:2]
+                frame_buffer.append(frame)
+        except Exception as e:
+            print(f"Error processing {img_path}: {e}")
+            continue
+    
+    if not frame_buffer:
+        raise ValueError("No valid frames found to create video")
+    
+    # 生成视频文件名
+    output_video_path = os.path.join(result_folder, f"manipulation.mp4")
+    
+    # 创建视频
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 或使用 'avc1' 更好的兼容性
+    video_writer = cv2.VideoWriter(
+        output_video_path,
+        fourcc,
+        fps,
+        (width, height)
+    )
+    
+    for frame in frame_buffer:
+        video_writer.write(frame)
+    
+    video_writer.release()
+    
+    # 验证视频是否成功创建
+    if not os.path.exists(output_video_path) or os.path.getsize(output_video_path) == 0:
+        raise RuntimeError(f"Failed to create video at {output_video_path}")
+    
+    return output_video_path
+
+
 def run_simulation(
     scene: str,
     prompt: str,
@@ -488,14 +562,16 @@ def run_simulation(
     status = get_task_status(task_id)
     print("status: ", status)
     if status.get("status") == "completed":
-        # time.sleep(3)
+        time.sleep(3)
         oss_video_path = os.path.join(result_folder, "manipulation.mp4")
         local_video_path = os.path.join(user_dir, task_id, "tasks", "manipulation.mp4")
-        download_oss_file(oss_video_path, local_video_path)
+        # download_oss_file(oss_video_path, local_video_path)
         print("oss_video_path: ", oss_video_path)
         print("local_video_path: ", local_video_path)
 
-        video_path = convert_to_h264(local_video_path)
+        video_path = generate_whole_video(task_id, request)
+        
+        # video_path = convert_to_h264(local_video_path)
 
         # 创建新的历史记录条目
         new_entry = {

requirements.txt

@@ -4,4 +4,5 @@ jsonlib-python3
 opencv-python
 numpy
 python-dateutil
-oss2
+oss2
+natsort