dev streaming with oss

app.py

@@ -11,6 +11,7 @@ import numpy as np
 from datetime import datetime, timedelta
 from collections import defaultdict
 import shutil
+from urllib.parse import urljoin
 
 # os.environ["SPACES_QUEUE_ENABLED"] = "true"
 
@@ -159,23 +160,57 @@ def format_logs_for_display(logs: list) -> str:
 
 
 ###############################################################################
+def list_public_oss_files(base_url: str) -> List[str]:
+    """列出公共OSS文件夹中的所有图片文件"""
+    # 注意：这需要OSS支持目录列表功能，或者你有预先知道的文件命名规则
+    # 如果OSS不支持目录列表，可能需要后端API提供文件列表
+    # 这里假设可以直接通过HTTP访问
+    
+    # 实际情况可能需要根据你的OSS具体配置调整
+    # 这里只是一个示例实现
+    try:
+        response = requests.get(base_url)
+        if response.status_code == 200:
+            # 这里需要根据OSS返回的实际内容解析文件列表
+            # 可能需要使用HTML解析器或正则表达式
+            # 以下只是示例
+            import re
+            files = re.findall(r'href="([^"]+\.(?:jpg|png|jpeg))"', response.text)
+            return sorted([urljoin(base_url, f) for f in files])
+        return []
+    except Exception as e:
+        print(f"Error listing public OSS files: {e}")
+        return []
 
+def download_public_file(url: str, local_path: str):
+    """下载公开可访问的文件"""
+    try:
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(local_path, 'wb') as f:
+                for chunk in response.iter_content(1024):
+                    f.write(chunk)
+            return True
+        return False
+    except Exception as e:
+        print(f"Error downloading public file {url}: {e}")
+        return False
 
-def stream_simulation_results(result_folder: str, task_id: str, fps: int = 30):
+def stream_simulation_results(result_folder: str, task_id: str, fps: int = 30, request: gr.Request):
     """
-    流式输出仿真结果，同时监控图片文件夹和后端任务状态
+    流式输出仿真结果，从公共OSS读取图片
     
     参数:
-        result_folder: 包含生成图片的文件夹路径
+        result_folder: OSS上包含生成图片的文件夹URL (从后端API返回)
         task_id: 后端任务ID用于状态查询
         fps: 输出视频的帧率
+        request: Gradio请求对象
         
     生成:
         生成的视频文件路径 (分段输出)
     """
     # 初始化变量
-    result_folder = os.path.join(result_folder, "image")
-    os.makedirs(result_folder, exist_ok=True)
+    image_folder = urljoin(result_folder, "image/")  # 确保以/结尾
     frame_buffer: List[np.ndarray] = []
     frames_per_segment = fps * 2  # 每2秒60帧
     processed_files = set()
@@ -184,6 +219,11 @@ def stream_simulation_results(result_folder: str, task_id: str, fps: int = 30):
     status_check_interval = 5  # 每5秒检查一次后端状态
     max_time = 240
 
+    # 创建临时目录存储下载的图片
+    user_dir = os.path.join(TMP_ROOT, str(request.session_hash))
+    local_image_dir = os.path.join(user_dir, "tasks", "images")
+    os.makedirs(local_image_dir, exist_ok=True)
+
     while max_time > 0:
         max_time -= 1
         current_time = time.time()
@@ -194,9 +234,9 @@ def stream_simulation_results(result_folder: str, task_id: str, fps: int = 30):
             print("status: ", status)
             if status.get("status") == "completed":
                 # 确保处理完所有已生成的图片
-                process_remaining_images(result_folder, processed_files, frame_buffer)
+                process_remaining_public_oss_images(image_folder, local_image_dir, processed_files, frame_buffer)
                 if frame_buffer:
-                    yield create_video_segment(frame_buffer, fps, width, height)
+                    yield create_video_segment(frame_buffer, fps, width, height, request)
                 break
             elif status.get("status") == "failed":
                 raise gr.Error(f"任务执行失败: {status.get('result', '未知错误')}")
@@ -204,35 +244,40 @@ def stream_simulation_results(result_folder: str, task_id: str, fps: int = 30):
                 break
             last_status_check = current_time
 
-        # 处理新生成的图片
-        current_files = sorted(
-            [f for f in os.listdir(result_folder) 
-             if f.lower().endswith(('.png', '.jpg', '.jpeg'))],
-            key=lambda x: os.path.splitext(x)[0]  # 按文件名排序
-        )
-        
-        new_files = [f for f in current_files if f not in processed_files]
-        has_new_frames = False
-        
-        for filename in new_files:
-            try:
-                img_path = os.path.join(result_folder, filename)
-                frame = cv2.imread(img_path)
-                if frame is not None:
-                    if width == 0:  # 第一次获取图像尺寸
-                        height, width = frame.shape[:2]
-                    
-                    frame_buffer.append(frame)
-                    processed_files.add(filename)
-                    has_new_frames = True
-            except Exception as e:
-                print(f"Error processing {filename}: {e}")
+        # 从公共OSS获取文件列表
+        try:
+            # 注意：这里假设可以直接列出OSS文件
+            # 如果不行，可能需要后端API提供文件列表
+            oss_files = list_public_oss_files(image_folder)
+            new_files = [f for f in oss_files if f not in processed_files]
+            has_new_frames = False
+            
+            for file_url in new_files:
+                try:
+                    # 下载文件到本地
+                    filename = os.path.basename(file_url)
+                    local_path = os.path.join(local_image_dir, filename)
+                    if download_public_file(file_url, local_path):
+                        # 读取图片
+                        frame = cv2.imread(local_path)
+                        if frame is not None:
+                            if width == 0:  # 第一次获取图像尺寸
+                                height, width = frame.shape[:2]
+                            
+                            frame_buffer.append(frame)
+                            processed_files.add(file_url)
+                            has_new_frames = True
+                except Exception as e:
+                    print(f"Error processing {file_url}: {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 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)
+        except Exception as e:
+            print(f"Error accessing public OSS: {e}")
 
         time.sleep(1)  # 避免过于频繁检查
     
@@ -255,25 +300,27 @@ def create_video_segment(frames: List[np.ndarray], fps: int, width: int, height:
     
     return segment_name
 
-def process_remaining_images(result_folder: str, processed_files: set, frame_buffer: List[np.ndarray]):
-    """处理剩余的图片"""
-    current_files = sorted(
-        [f for f in os.listdir(result_folder) 
-         if f.lower().endswith(('.png', '.jpg', '.jpeg'))],
-        key=lambda x: os.path.splitext(x)[0]
-    )
-    
-    new_files = [f for f in current_files if f not in processed_files]
-    
-    for filename in new_files:
-        try:
-            img_path = os.path.join(result_folder, filename)
-            frame = cv2.imread(img_path)
-            if frame is not None:
-                frame_buffer.append(frame)
-                processed_files.add(filename)
-        except Exception as e:
-            print(f"Error processing remaining {filename}: {e}")
+def process_remaining_public_oss_images(oss_folder: str, local_dir: str, processed_files: set, frame_buffer: List[np.ndarray]):
+    """处理公共OSS上剩余的图片"""
+    try:
+        oss_files = list_public_oss_files(oss_folder)
+        new_files = [f for f in oss_files if f not in processed_files]
+        
+        for file_url in new_files:
+            try:
+                # 下载文件到本地
+                filename = os.path.basename(file_url)
+                local_path = os.path.join(local_dir, filename)
+                if download_public_file(file_url, local_path):
+                    # 读取图片
+                    frame = cv2.imread(local_path)
+                    if frame is not None:
+                        frame_buffer.append(frame)
+                        processed_files.add(file_url)
+            except Exception as e:
+                print(f"Error processing remaining {file_url}: {e}")
+    except Exception as e:
+        print(f"Error accessing public OSS for remaining files: {e}")