294 lines
9.8 KiB
Python
294 lines
9.8 KiB
Python
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) |