import torch import cv2 import os import glob import numpy as np import onnxruntime as ort from tqdm import tqdm from concurrent.futures import ThreadPoolExecutor import threading def preview_frame(video_path, frame_number): """预览指定帧的图像""" cap = cv2.VideoCapture(video_path) cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number) ret, frame = cap.read() cap.release() if ret: cv2.imwrite(f'frame_{frame_number}.jpg', frame) print(f"已保存第 {frame_number} 帧的预览图像") else: print(f"无法读取第 {frame_number} 帧") def extract_specific_frames(video_path, frame_numbers): """从视频中提取指定帧号的图像""" cap = cv2.VideoCapture(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)) print(f"视频信息：\n分辨率: {frame_width}x{frame_height}\n总帧数: {total_frames}") # 创建画布 - 宽度为单帧宽度乘以要提取的帧数 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): print(f"正在处理第 {target_frame} 帧...") # 快进到目标帧 while current_frame < target_frame: ret = cap.grab() # 使用grab()而不是read()来快速跳过帧 if not ret: print(f"无法读取到第 {target_frame} 帧") break current_frame += 1 # 读取目标帧 ret, frame = cap.retrieve() # 获取当前帧的图像数据 if ret: # 将帧放入画布 canvas[:, canvas_position:canvas_position+frame_width] = frame canvas_position += frame_width # 保存每个提取的帧（可选） cv2.imwrite(f'frame_{target_frame}.jpg', frame) else: print(f"无法读取第 {target_frame} 帧的图像数据") cap.release() # 如果canvas_position小于预期的canvas_width，裁剪未使用的部分 if canvas_position < canvas_width: canvas = canvas[:, :canvas_position] return canvas def preview_all_frames(video_path, interval=10): """每隔一定间隔预览并保存帧，用于选择合适的帧""" cap = cv2.VideoCapture(video_path) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) frame_count = 0 while True: ret, frame = cap.read() if not ret: break if frame_count % interval == 0: cv2.imwrite(f'preview_frame_{frame_count}.jpg', frame) print(f"已保存第 {frame_count} 帧的预览图像") frame_count += 1 cap.release() print(f"总共处理了 {frame_count} 帧") def get_longimage(video_path2,output_path): video_path = video_path2 # 是否需要预览帧 need_preview = False if need_preview: # 预览特定帧 preview_frames = [13, 82, 196, 290, 381] for frame_num in preview_frames: preview_frame(video_path, frame_num) # 或者每隔10帧预览一次 # preview_all_frames(video_path, interval=10) return # 指定要提取的帧号 frame_numbers = [13, 82, 196, 290, 381, 483, 697, 798, 893, 1047] # 提取帧并拼接 panorama = extract_specific_frames(video_path, frame_numbers) # 保存结果 cv2.imwrite(os.path.join(output_path,'panorama.jpg'), panorama) print(f"全景图已保存，尺寸: {panorama.shape[1]}x{panorama.shape[0]}") def init_onnx(): """Initialize ONNX Runtime session""" test_weights = "/data1/Hatcher/model/PopulationCounting/Bayesian_project_ours/sim-CAS-base-test.onnx" if 'CUDAExecutionProvider' in ort.get_available_providers(): print("CUDA is available. Using GPU.") providers = ["CUDAExecutionProvider"] if ort.get_device() == "GPU" else ["CPUExecutionProvider"] else: print("CUDA not available. Using CPU.") providers = ["CPUExecutionProvider"] ort_session = ort.InferenceSession(test_weights, providers=providers) return ort_session def split_image(image): """Split large image into 800x800 patches""" h, w = image.shape[:2] # Calculate optimal width for resizing n_segments = max(1, w // 800) target_w = n_segments * 800 scale = target_w / w # Resize image to make width divisible by 800 new_w = target_w new_h = int(h * scale) resized_img = cv2.resize(image, (new_w, new_h)) # Split into patches patches = [] positions = [] # Store original positions for reconstruction for i in range(0, new_w, 800): patch = resized_img[:, i:i+800] if patch.shape[1] < 800: # Handle last patch if smaller than 800 continue patches.append(patch) positions.append((i, i+800)) return patches, positions, scale, (new_h, new_w) def process_patch(patch, ort_session): """Process a single 800x800 patch""" # Preprocess 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) # Run inference 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], input_feed={onnx_input_name: input_blob}) # Get density map density_map = onnx_result[0][0][0] density_map[density_map<1e-6] = 0 return density_map def process_single_image(image_path, save_root, ort_session, file_lock): """Process a single large image""" try: # Get image name image_name = os.path.basename(image_path) result_name = f"result_{image_name}" # Read image img_or = cv2.imread(image_path) if img_or is None: print(f"Failed to read image: {image_path}") return False # Split image into patches patches, positions, scale, (new_h, new_w) = split_image(img_or) # Process each patch 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) # Combine density maps 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 # Resize combined density map back to original size combined_density = cv2.resize(combined_density, (img_or.shape[1], img_or.shape[0])) # Create heatmap heatmap = combined_density/np.max(combined_density) heatmap = np.power(heatmap, 0.7) # 降低gamma值会使颜色更深 heatmap = np.uint8(heatmap*255) heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) # Create visualization 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)) # Add density map to top-left corner img_or[0:imgheatmap_H, 0:imgheatmap_W] = density_result # Add count text cv2.putText(img_or, f"Count: {int(total_count)}", (50,200), cv2.FONT_HERSHEY_SIMPLEX, 4, (0,0,255), 8) # Save result save_path = os.path.join(save_root, result_name) with file_lock: cv2.imwrite(save_path, img_or) print(save_path) return True except Exception as e: print(f"Error processing {image_path}: {str(e)}") return False def process_images_batch(input_folder, num_threads=4): """Process images in parallel""" # Initialize ONNX session ort_session = init_onnx() # Create output directory save_root = os.path.join(input_folder, "results") os.makedirs(save_root, exist_ok=True) # Get image paths image_paths = glob.glob(os.path.join(input_folder, "*.jpg")) + \ glob.glob(os.path.join(input_folder, "*.png")) print(f"Found {len(image_paths)} images to process") # Initialize counters and locks success_count = 0 counter_lock = threading.Lock() file_lock = threading.Lock() # Create progress bar pbar = tqdm(total=len(image_paths), desc="Processing images") def process_with_progress(image_path): nonlocal success_count result = process_single_image(image_path, save_root, ort_session, file_lock) with counter_lock: if result: success_count += 1 pbar.update(1) pbar.set_postfix({'succeeded': success_count}) # Process images in parallel with ThreadPoolExecutor(max_workers=num_threads) as executor: executor.map(process_with_progress, image_paths) pbar.close() print(f"Processing complete. Successfully processed {success_count} images") if __name__ == "__main__": video_path = "/data1/Hatcher/model/project/MPKFNet/mmexport1733810879708.mp4" out_put_path = "./20251105001" if not os.path.exists(out_put_path): os.makedirs(out_put_path) get_longimage(video_path,out_put_path) # Start processing process_images_batch(out_put_path)