176 lines
6.7 KiB
Python
176 lines
6.7 KiB
Python
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
|