mp4保存
This commit is contained in:
+103
-29
@@ -8,6 +8,29 @@ GPU-accelerated face detection using ONNX Runtime.
|
||||
import os
|
||||
import sys
|
||||
import platform
|
||||
|
||||
# CRITICAL: Fix LD_LIBRARY_PATH before importing cv2 or torch
|
||||
# cv2 adds its own lib path to the front, which can override ROCm libraries
|
||||
def fix_library_path():
|
||||
"""Ensure ROCm libraries are loaded before cv2's bundled libraries."""
|
||||
ld_path = os.environ.get('LD_LIBRARY_PATH', '')
|
||||
|
||||
# Split and filter paths
|
||||
paths = [p for p in ld_path.split(':') if p]
|
||||
|
||||
# Separate ROCm/GPU paths from other paths
|
||||
rocm_paths = [p for p in paths if 'rocm' in p.lower() or 'clr-' in p or 'hip' in p.lower()]
|
||||
other_paths = [p for p in paths if p not in rocm_paths]
|
||||
|
||||
# Rebuild with ROCm paths first
|
||||
if rocm_paths:
|
||||
new_ld_path = ':'.join(rocm_paths + other_paths)
|
||||
os.environ['LD_LIBRARY_PATH'] = new_ld_path
|
||||
print(f"[FaceMask] Fixed LD_LIBRARY_PATH to prioritize ROCm libraries")
|
||||
|
||||
# Fix library path BEFORE any other imports
|
||||
fix_library_path()
|
||||
|
||||
import threading
|
||||
import uuid
|
||||
import queue
|
||||
@@ -61,11 +84,15 @@ class GenerateRequest(BaseModel):
|
||||
mask_scale: float = 1.5
|
||||
|
||||
def process_video_task(task_id: str, req: GenerateRequest):
|
||||
"""Background task to process video."""
|
||||
"""Background task to process video with async MP4 output."""
|
||||
writer = None
|
||||
write_queue = None
|
||||
writer_thread = None
|
||||
|
||||
try:
|
||||
tasks[task_id].status = TaskStatus.PROCESSING
|
||||
cancel_event = cancel_events.get(task_id)
|
||||
|
||||
|
||||
# Verify video exists
|
||||
if not os.path.exists(req.video_path):
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
@@ -78,28 +105,60 @@ def process_video_task(task_id: str, req: GenerateRequest):
|
||||
conf_threshold=req.conf_threshold,
|
||||
iou_threshold=req.iou_threshold
|
||||
)
|
||||
# Ensure model is loaded
|
||||
_ = detector.model
|
||||
|
||||
|
||||
# Open video
|
||||
cap = cv2.VideoCapture(req.video_path)
|
||||
if not cap.isOpened():
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Failed to open video"
|
||||
return
|
||||
|
||||
# Determine frame range
|
||||
|
||||
# Get video properties
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
total_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
end_frame = min(req.end_frame, total_video_frames - 1)
|
||||
frames_to_process = end_frame - req.start_frame + 1
|
||||
|
||||
|
||||
tasks[task_id].total = frames_to_process
|
||||
|
||||
|
||||
# Ensure output directory exists
|
||||
os.makedirs(req.output_dir, exist_ok=True)
|
||||
|
||||
print(f"Starting processing: {req.video_path} ({frames_to_process} frames)")
|
||||
|
||||
|
||||
# Setup MP4 writer (grayscale)
|
||||
output_video_path = os.path.join(req.output_dir, "mask.mp4")
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height), isColor=False)
|
||||
|
||||
if not writer.isOpened():
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Failed to create video writer"
|
||||
cap.release()
|
||||
return
|
||||
|
||||
# Async writer setup
|
||||
write_queue = queue.Queue(maxsize=30) # Buffer up to 30 frames
|
||||
writer_running = threading.Event()
|
||||
writer_running.set()
|
||||
|
||||
def async_writer():
|
||||
"""Background thread for writing frames to video."""
|
||||
while writer_running.is_set() or not write_queue.empty():
|
||||
try:
|
||||
mask = write_queue.get(timeout=0.1)
|
||||
if mask is not None:
|
||||
writer.write(mask)
|
||||
write_queue.task_done()
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
writer_thread = threading.Thread(target=async_writer, daemon=True)
|
||||
writer_thread.start()
|
||||
|
||||
print(f"Starting processing: {req.video_path} ({frames_to_process} frames) -> {output_video_path}")
|
||||
|
||||
# Process loop
|
||||
current_count = 0
|
||||
for frame_idx in range(req.start_frame, end_frame + 1):
|
||||
@@ -107,38 +166,44 @@ def process_video_task(task_id: str, req: GenerateRequest):
|
||||
tasks[task_id].status = TaskStatus.CANCELLED
|
||||
tasks[task_id].message = "Cancelled by user"
|
||||
break
|
||||
|
||||
|
||||
# Read frame
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
|
||||
ret, frame = cap.read()
|
||||
|
||||
|
||||
if ret:
|
||||
# Detect
|
||||
detections = detector.detect(frame)
|
||||
|
||||
|
||||
# Generate mask
|
||||
mask = detector.generate_mask(
|
||||
frame.shape,
|
||||
detections,
|
||||
mask_scale=req.mask_scale
|
||||
)
|
||||
|
||||
# Save
|
||||
mask_filename = f"mask_{current_count:06d}.png" # Note: using relative index for filename
|
||||
mask_path = os.path.join(req.output_dir, mask_filename)
|
||||
cv2.imwrite(mask_path, mask)
|
||||
|
||||
|
||||
# Async write to queue
|
||||
write_queue.put(mask)
|
||||
|
||||
# Update progress
|
||||
current_count += 1
|
||||
tasks[task_id].progress = current_count
|
||||
|
||||
|
||||
# Cleanup
|
||||
writer_running.clear()
|
||||
write_queue.join() # Wait for all frames to be written
|
||||
if writer_thread:
|
||||
writer_thread.join(timeout=5)
|
||||
|
||||
cap.release()
|
||||
|
||||
if writer:
|
||||
writer.release()
|
||||
|
||||
if tasks[task_id].status == TaskStatus.PROCESSING:
|
||||
tasks[task_id].status = TaskStatus.COMPLETED
|
||||
tasks[task_id].result_path = req.output_dir
|
||||
tasks[task_id].result_path = output_video_path # Return video path
|
||||
tasks[task_id].message = "Processing completed successfully"
|
||||
print(f"Task {task_id} completed.")
|
||||
print(f"Task {task_id} completed: {output_video_path}")
|
||||
|
||||
except Exception as e:
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
@@ -230,12 +295,21 @@ def log_startup_diagnostics():
|
||||
for var in rocm_vars:
|
||||
value = os.environ.get(var)
|
||||
if value:
|
||||
# Truncate very long values
|
||||
if len(value) > 200:
|
||||
display_value = value[:200] + "... (truncated)"
|
||||
# For LD_LIBRARY_PATH, show if ROCm paths are included
|
||||
if var == 'LD_LIBRARY_PATH':
|
||||
has_rocm = 'rocm' in value.lower() or 'clr-' in value
|
||||
has_hip = 'hip' in value.lower()
|
||||
print(f" {var}: {value[:100]}...")
|
||||
print(f" Contains ROCm paths: {has_rocm}")
|
||||
print(f" Contains HIP paths: {has_hip}")
|
||||
if not has_rocm:
|
||||
print(f" ⚠️ WARNING: ROCm library paths not found!")
|
||||
else:
|
||||
display_value = value
|
||||
print(f" {var}: {display_value}")
|
||||
if len(value) > 200:
|
||||
display_value = value[:200] + "... (truncated)"
|
||||
else:
|
||||
display_value = value
|
||||
print(f" {var}: {display_value}")
|
||||
else:
|
||||
print(f" {var}: (not set)")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user