first commit

sys_long_crowd_jiekou.py

@@ -0,0 +1,294 @@
+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)