import os import cv2 import numpy as np import onnxruntime as ort from concurrent.futures import ThreadPoolExecutor import threading from typing import List, Tuple # 全局 ONNX 会话（单例） _ort_session = None _session_lock = threading.Lock() def get_ort_session(): """获取 ONNX Runtime 会话（线程安全单例）""" global _ort_session if _ort_session is None: with _session_lock: if _ort_session is None: test_weights = "/data1/Hatcher/model/PopulationCounting/Bayesian_project_ours/sim-CAS-base-test.onnx" if not os.path.exists(test_weights): raise FileNotFoundError(f"ONNX model not found: {test_weights}") providers = ["CUDAExecutionProvider"] if "CUDAExecutionProvider" in ort.get_available_providers() else ["CPUExecutionProvider"] _ort_session = ort.InferenceSession(test_weights, providers=providers) print(f"ONNX model loaded with providers: {providers}") return _ort_session def extract_specific_frames(video_path: str, frame_numbers: List[int]) -> np.ndarray: """从视频中提取指定帧号的图像并拼接""" cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError(f"Cannot open video file: {video_path}") frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) canvas_width = frame_width * len(frame_numbers) canvas = np.zeros((frame_height, canvas_width, 3), dtype=np.uint8) current_frame = 0 canvas_position = 0 for target_frame in sorted(frame_numbers): if target_frame >= total_frames: continue while current_frame < target_frame: if not cap.grab(): break current_frame += 1 ret, frame = cap.retrieve() if ret: canvas[:, canvas_position:canvas_position+frame_width] = frame canvas_position += frame_width else: print(f"Warning: Failed to read frame {target_frame}") cap.release() if canvas_position < canvas_width: canvas = canvas[:, :canvas_position] return canvas def split_image(image: np.ndarray) -> Tuple[List[np.ndarray], List[Tuple[int, int]], float, Tuple[int, int]]: """Split large image into 800x800 patches""" h, w = image.shape[:2] n_segments = max(1, w // 800) target_w = n_segments * 800 scale = target_w / w new_w = target_w new_h = int(h * scale) resized_img = cv2.resize(image, (new_w, new_h)) patches = [] positions = [] for i in range(0, new_w, 800): patch = resized_img[:, i:i+800] if patch.shape[1] == 800: patches.append(patch) positions.append((i, i+800)) return patches, positions, scale, (new_h, new_w) def process_patch(patch: np.ndarray, ort_session) -> np.ndarray: """Process a single 800x800 patch""" img = cv2.resize(patch, (800, 800)) / 255.0 img = img[np.newaxis, :, :].transpose(0, 3, 1, 2) input_blob = np.array(img, dtype=np.float32) onnx_input_name = ort_session.get_inputs()[0].name onnx_outputs_names = ort_session.get_outputs()[0].name onnx_result = ort_session.run([onnx_outputs_names], {onnx_input_name: input_blob}) density_map = onnx_result[0][0][0] density_map[density_map < 1e-6] = 0 return density_map def process_single_image(image_path: str, save_root: str, ort_session, file_lock) -> bool: """Process a single large image""" try: image_name = os.path.basename(image_path) result_name = f"result_{image_name}" img_or = cv2.imread(image_path) if img_or is None: return False patches, positions, scale, (new_h, new_w) = split_image(img_or) density_maps = [] total_count = 0 for patch in patches: density_map = process_patch(patch, ort_session) density_maps.append(density_map) total_count += np.sum(density_map) combined_density = np.zeros((new_h, new_w)) for density_map, (start_x, end_x) in zip(density_maps, positions): resized_density = cv2.resize(density_map, (800, new_h)) combined_density[:, start_x:end_x] = resized_density combined_density = cv2.resize(combined_density, (img_or.shape[1], img_or.shape[0])) heatmap = combined_density / np.max(combined_density) heatmap = np.power(heatmap, 0.7) heatmap = np.uint8(heatmap * 255) heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) imgheatmap_W = int(img_or.shape[1] / 3) imgheatmap_H = int(img_or.shape[0] / 3) density_result = cv2.resize(heatmap, (imgheatmap_W, imgheatmap_H)) img_or[0:imgheatmap_H, 0:imgheatmap_W] = density_result cv2.putText(img_or, f"Count: {int(total_count)}", (50, 200), cv2.FONT_HERSHEY_SIMPLEX, 4, (0, 0, 255), 8) save_path = os.path.join(save_root, result_name) with file_lock: cv2.imwrite(save_path, img_or) return True, float(total_count) except Exception as e: print(f"Error processing {image_path}: {str(e)}") return False def process_video_to_panorama(video_path: str, output_dir: str) -> Tuple[str, float]: """Main pipeline: video -> panorama -> density estimation, returns (result_path, total_count)""" os.makedirs(output_dir, exist_ok=True) # Step 1: Extract frames and create panorama frame_numbers = [13, 82, 196, 290, 381, 483, 697, 798, 893, 1047] panorama = extract_specific_frames(video_path, frame_numbers) panorama_path = os.path.join(output_dir, "panorama.jpg") cv2.imwrite(panorama_path, panorama) # Step 2: Process images with ONNX model ort_session = get_ort_session() image_paths = [panorama_path] save_root = output_dir os.makedirs(save_root, exist_ok=True) results = [] file_lock = threading.Lock() def process_with_result(image_path): result = process_single_image(image_path, save_root, ort_session, file_lock) results.append(result) # (success, total_count) with ThreadPoolExecutor(max_workers=4) as executor: executor.map(process_with_result, image_paths) if not results or not results[0][0]: raise RuntimeError("Failed to process panorama image") result_path = os.path.join(save_root, f"result_{os.path.basename(panorama_path)}") total_count = results[0][1] # ← 获取 total_count return result_path, total_count # ← 返回 total_count