first commit

app/processor.py

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