917 lines
31 KiB
Python
917 lines
31 KiB
Python
"""
|
||
Face Detection Inference Server.
|
||
|
||
This FastAPI application runs in a separate process to handle
|
||
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("[FaceMask] Fixed LD_LIBRARY_PATH to prioritize ROCm libraries")
|
||
|
||
# Fix library path BEFORE any other imports
|
||
fix_library_path()
|
||
|
||
import queue
|
||
import threading
|
||
import uuid
|
||
import traceback
|
||
import subprocess
|
||
from typing import Dict, Optional, List
|
||
from pathlib import Path
|
||
|
||
from fastapi import FastAPI, HTTPException, BackgroundTasks
|
||
from pydantic import BaseModel
|
||
import uvicorn
|
||
import cv2
|
||
import numpy as np
|
||
import msgpack
|
||
|
||
# Add project root to path for imports if needed
|
||
sys.path.append(str(Path(__file__).parent.parent))
|
||
|
||
from server.detector import get_detector
|
||
|
||
app = FastAPI(title="Face Mask Inference Server")
|
||
|
||
# GPU status cache
|
||
_gpu_status_cache = None
|
||
|
||
# Task storage
|
||
class TaskStatus:
|
||
PENDING = "pending"
|
||
PROCESSING = "processing"
|
||
COMPLETED = "completed"
|
||
FAILED = "failed"
|
||
CANCELLED = "cancelled"
|
||
|
||
class Task(BaseModel):
|
||
id: str
|
||
status: str
|
||
progress: int = 0
|
||
total: int = 0
|
||
message: Optional[str] = None
|
||
result_path: Optional[str] = None
|
||
|
||
# In-memory storage
|
||
tasks: Dict[str, Task] = {}
|
||
cancel_events: Dict[str, threading.Event] = {}
|
||
|
||
class GenerateRequest(BaseModel):
|
||
video_path: str
|
||
output_dir: str
|
||
start_frame: int
|
||
end_frame: int
|
||
conf_threshold: float = 0.5
|
||
iou_threshold: float = 0.45
|
||
|
||
|
||
class VideoInfoRequest(BaseModel):
|
||
video_path: str
|
||
|
||
|
||
class BakeRequest(BaseModel):
|
||
video_path: str
|
||
detections_path: str
|
||
output_path: str
|
||
blur_size: int = 50
|
||
display_scale: float = 1.0
|
||
format: str = "mp4"
|
||
|
||
|
||
class _FFmpegPipeWriter:
|
||
"""Write BGR frames to ffmpeg stdin."""
|
||
|
||
def __init__(self, cmd: List[str]):
|
||
self._proc = subprocess.Popen(
|
||
cmd,
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=subprocess.PIPE,
|
||
)
|
||
if self._proc.stdin is None:
|
||
self._proc.kill()
|
||
raise RuntimeError("Failed to open ffmpeg stdin")
|
||
|
||
def write(self, frame: np.ndarray) -> None:
|
||
if self._proc.stdin is None:
|
||
raise RuntimeError("ffmpeg stdin is not available")
|
||
self._proc.stdin.write(frame.tobytes())
|
||
|
||
def release(self) -> None:
|
||
if self._proc.stdin:
|
||
try:
|
||
self._proc.stdin.close()
|
||
except Exception:
|
||
pass
|
||
|
||
stderr_output = ""
|
||
if self._proc.stderr:
|
||
try:
|
||
stderr_output = self._proc.stderr.read().decode("utf-8", errors="replace")
|
||
except Exception:
|
||
stderr_output = ""
|
||
|
||
rc = self._proc.wait()
|
||
if rc != 0:
|
||
tail = "\n".join(stderr_output.strip().splitlines()[-8:])
|
||
raise RuntimeError(f"ffmpeg writer failed (code={rc}): {tail}")
|
||
|
||
|
||
def _build_ffmpeg_vaapi_writer(
|
||
output_path: str,
|
||
fps: float,
|
||
width: int,
|
||
height: int,
|
||
) -> _FFmpegPipeWriter:
|
||
"""Create ffmpeg h264_vaapi writer with QP=24 (balanced quality/speed)."""
|
||
cmd = [
|
||
"ffmpeg",
|
||
"-hide_banner",
|
||
"-loglevel",
|
||
"error",
|
||
"-y",
|
||
"-vaapi_device",
|
||
"/dev/dri/renderD128",
|
||
"-f",
|
||
"rawvideo",
|
||
"-pix_fmt",
|
||
"bgr24",
|
||
"-s",
|
||
f"{width}x{height}",
|
||
"-r",
|
||
f"{fps}",
|
||
"-i",
|
||
"-",
|
||
"-an",
|
||
"-vf",
|
||
"format=nv12,hwupload",
|
||
"-c:v",
|
||
"h264_vaapi",
|
||
"-qp",
|
||
"24",
|
||
output_path,
|
||
]
|
||
return _FFmpegPipeWriter(cmd)
|
||
|
||
|
||
def _build_video_writer(
|
||
output_path: str,
|
||
fmt: str,
|
||
fps: float,
|
||
width: int,
|
||
height: int,
|
||
) -> object:
|
||
"""Create writer with VAAPI preference and OpenCV fallback."""
|
||
format_key = fmt.lower()
|
||
|
||
if format_key in {"mp4", "mov"}:
|
||
try:
|
||
writer = _build_ffmpeg_vaapi_writer(output_path, fps, width, height)
|
||
print("[FaceMask] Using output encoder: ffmpeg h264_vaapi (-qp 24)")
|
||
return writer
|
||
except Exception as e:
|
||
print(f"[FaceMask] VAAPI writer unavailable, fallback to OpenCV: {e}")
|
||
|
||
codec_candidates = {
|
||
"mp4": ["avc1", "mp4v"],
|
||
"mov": ["avc1", "mp4v"],
|
||
"avi": ["MJPG", "XVID"],
|
||
}.get(format_key, ["mp4v"])
|
||
|
||
for codec in codec_candidates:
|
||
writer = cv2.VideoWriter(
|
||
output_path,
|
||
cv2.VideoWriter_fourcc(*codec),
|
||
fps,
|
||
(width, height),
|
||
isColor=True,
|
||
)
|
||
if writer.isOpened():
|
||
print(f"[FaceMask] Using output codec: {codec}")
|
||
return writer
|
||
writer.release()
|
||
|
||
raise RuntimeError(f"Failed to create video writer for format='{fmt}'")
|
||
|
||
def _scale_bbox(
|
||
x: int,
|
||
y: int,
|
||
w: int,
|
||
h: int,
|
||
scale: float,
|
||
frame_width: int,
|
||
frame_height: int,
|
||
) -> Optional[List[int]]:
|
||
"""Scale bbox around center and clamp to frame boundaries."""
|
||
if w <= 0 or h <= 0:
|
||
return None
|
||
|
||
center_x = x + (w * 0.5)
|
||
center_y = y + (h * 0.5)
|
||
scaled_w = max(1, int(w * scale))
|
||
scaled_h = max(1, int(h * scale))
|
||
|
||
x1 = max(0, int(center_x - scaled_w * 0.5))
|
||
y1 = max(0, int(center_y - scaled_h * 0.5))
|
||
x2 = min(frame_width, x1 + scaled_w)
|
||
y2 = min(frame_height, y1 + scaled_h)
|
||
out_w = x2 - x1
|
||
out_h = y2 - y1
|
||
if out_w <= 0 or out_h <= 0:
|
||
return None
|
||
return [x1, y1, out_w, out_h]
|
||
|
||
|
||
def process_video_task(task_id: str, req: GenerateRequest):
|
||
"""Background task to detect faces and save bbox cache as msgpack."""
|
||
cap = None
|
||
|
||
try:
|
||
tasks[task_id].status = TaskStatus.PROCESSING
|
||
cancel_event = cancel_events.get(task_id)
|
||
|
||
if not os.path.exists(req.video_path):
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = f"Video not found: {req.video_path}"
|
||
return
|
||
|
||
print(f"Loading detector for task {task_id}...")
|
||
detector = get_detector(
|
||
conf_threshold=req.conf_threshold,
|
||
iou_threshold=req.iou_threshold,
|
||
)
|
||
_ = detector.model
|
||
|
||
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
|
||
|
||
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||
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
|
||
if frames_to_process <= 0:
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = "Invalid frame range"
|
||
return
|
||
|
||
tasks[task_id].total = frames_to_process
|
||
os.makedirs(req.output_dir, exist_ok=True)
|
||
output_msgpack_path = os.path.join(req.output_dir, "detections.msgpack")
|
||
|
||
if req.start_frame > 0:
|
||
seek_ok = cap.set(cv2.CAP_PROP_POS_FRAMES, req.start_frame)
|
||
if not seek_ok:
|
||
for _ in range(req.start_frame):
|
||
ret, _ = cap.read()
|
||
if not ret:
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = (
|
||
f"Failed to seek to start frame: {req.start_frame}"
|
||
)
|
||
return
|
||
|
||
frame_buffer: List[np.ndarray] = []
|
||
frame_detections: List[List[List[float]]] = []
|
||
batch_size = 5
|
||
current_count = 0
|
||
|
||
def process_batch():
|
||
nonlocal current_count
|
||
if not frame_buffer:
|
||
return
|
||
|
||
batch_detections = detector.detect_batch(frame_buffer)
|
||
for detections in batch_detections:
|
||
packed_detections: List[List[float]] = []
|
||
for x, y, w, h, conf in detections:
|
||
# bboxをそのまま保存（表示スケールはBake時に適用）
|
||
bx, by, bw, bh = int(x), int(y), int(w), int(h)
|
||
bx = max(0, bx)
|
||
by = max(0, by)
|
||
bw = min(width - bx, bw)
|
||
bh = min(height - by, bh)
|
||
if bw <= 0 or bh <= 0:
|
||
continue
|
||
packed_detections.append([bx, by, bw, bh, float(conf)])
|
||
frame_detections.append(packed_detections)
|
||
current_count += 1
|
||
tasks[task_id].progress = current_count
|
||
|
||
frame_buffer.clear()
|
||
|
||
print(
|
||
f"Starting detection cache generation: {req.video_path} "
|
||
f"({frames_to_process} frames) -> {output_msgpack_path}"
|
||
)
|
||
|
||
for _ in range(req.start_frame, end_frame + 1):
|
||
if cancel_event and cancel_event.is_set():
|
||
tasks[task_id].status = TaskStatus.CANCELLED
|
||
tasks[task_id].message = "Cancelled by user"
|
||
break
|
||
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
break
|
||
|
||
frame_buffer.append(frame)
|
||
if len(frame_buffer) >= batch_size:
|
||
process_batch()
|
||
|
||
if frame_buffer:
|
||
process_batch()
|
||
|
||
if tasks[task_id].status == TaskStatus.PROCESSING:
|
||
payload = {
|
||
"version": 1,
|
||
"video_path": req.video_path,
|
||
"start_frame": req.start_frame,
|
||
"end_frame": req.start_frame + len(frame_detections) - 1,
|
||
"width": width,
|
||
"height": height,
|
||
"fps": fps,
|
||
"mask_scale": 1.0,
|
||
"frames": frame_detections,
|
||
}
|
||
with open(output_msgpack_path, "wb") as f:
|
||
f.write(msgpack.packb(payload, use_bin_type=True))
|
||
|
||
tasks[task_id].status = TaskStatus.COMPLETED
|
||
tasks[task_id].result_path = output_msgpack_path
|
||
tasks[task_id].message = "Detection cache completed"
|
||
print(f"Task {task_id} completed: {output_msgpack_path}")
|
||
|
||
except Exception as e:
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = str(e)
|
||
print(f"Error in task {task_id}: {e}")
|
||
traceback.print_exc()
|
||
finally:
|
||
if cap:
|
||
cap.release()
|
||
if task_id in cancel_events:
|
||
del cancel_events[task_id]
|
||
|
||
|
||
def process_bake_task(task_id: str, req: BakeRequest):
|
||
"""Bake blur using async pipeline: read/process/write run in parallel for 1.35x speedup."""
|
||
try:
|
||
tasks[task_id].status = TaskStatus.PROCESSING
|
||
cancel_event = cancel_events.get(task_id)
|
||
|
||
if not os.path.exists(req.video_path):
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = f"Video not found: {req.video_path}"
|
||
return
|
||
if not os.path.exists(req.detections_path):
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = f"Detections file not found: {req.detections_path}"
|
||
return
|
||
|
||
with open(req.detections_path, "rb") as f:
|
||
payload = msgpack.unpackb(f.read(), raw=False)
|
||
frames_detections = payload.get("frames")
|
||
if not isinstance(frames_detections, list):
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = "Invalid detections format: 'frames' is missing"
|
||
return
|
||
|
||
# 検出キャッシュの開始フレーム（ソース動画のフレームインデックス）
|
||
det_start_frame = int(payload.get("start_frame", 0))
|
||
|
||
# Get video info
|
||
temp_cap = cv2.VideoCapture(req.video_path)
|
||
if not temp_cap.isOpened():
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = "Failed to open source video"
|
||
return
|
||
|
||
src_fps = temp_cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||
src_width = int(temp_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
src_height = int(temp_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
src_frames = int(temp_cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
temp_cap.release()
|
||
|
||
if src_width <= 0 or src_height <= 0:
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = "Invalid source video dimensions"
|
||
return
|
||
|
||
# ソース動画の全フレームを出力（スワップ後もトリム設定が正しく機能するよう）
|
||
total = src_frames if src_frames > 0 else (det_start_frame + len(frames_detections))
|
||
if total <= 0:
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = "Source/detections frame count is zero"
|
||
return
|
||
tasks[task_id].total = total
|
||
|
||
output_dir = os.path.dirname(req.output_path)
|
||
if output_dir:
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# Pipeline setup
|
||
blur_size = max(1, int(req.blur_size))
|
||
if blur_size % 2 == 0:
|
||
blur_size += 1
|
||
display_scale = max(0.1, float(req.display_scale))
|
||
# blur_margin は境界問題回避のための計算用余白のみ（表示には使わない）
|
||
blur_margin = blur_size // 2
|
||
|
||
# Queues
|
||
queue_size = 8
|
||
read_queue: queue.Queue = queue.Queue(maxsize=queue_size)
|
||
process_queue: queue.Queue = queue.Queue(maxsize=queue_size)
|
||
|
||
# Shared state
|
||
error_holder = {"error": None}
|
||
progress_lock = threading.Lock()
|
||
current_progress = [0]
|
||
|
||
def _reader_worker():
|
||
"""Read frames from video."""
|
||
import time as _time
|
||
cap = cv2.VideoCapture(req.video_path)
|
||
if not cap.isOpened():
|
||
error_holder["error"] = "Failed to open video in reader"
|
||
return
|
||
|
||
t_read_total = 0.0
|
||
frame_count = 0
|
||
try:
|
||
for idx in range(total):
|
||
if cancel_event and cancel_event.is_set():
|
||
break
|
||
|
||
t0 = _time.perf_counter()
|
||
ok, frame = cap.read()
|
||
t_read_total += _time.perf_counter() - t0
|
||
|
||
if not ok:
|
||
break
|
||
|
||
read_queue.put((idx, frame))
|
||
frame_count += 1
|
||
except Exception as e:
|
||
error_holder["error"] = f"Reader error: {e}"
|
||
finally:
|
||
cap.release()
|
||
read_queue.put(None) # Sentinel
|
||
if frame_count > 0:
|
||
print(
|
||
f"[Perf/Reader] FINAL frame={frame_count}"
|
||
f" read_avg={t_read_total/frame_count*1000:.1f}ms"
|
||
f" throughput≈{frame_count/max(t_read_total,1e-9):.1f}fps"
|
||
)
|
||
|
||
def _processor_worker():
|
||
"""Process frames with ROI blur."""
|
||
import time as _time
|
||
t_wait_total = 0.0
|
||
t_blur_total = 0.0
|
||
t_blend_total = 0.0
|
||
frame_count = 0
|
||
REPORT_INTERVAL = 50
|
||
|
||
try:
|
||
while True:
|
||
if cancel_event and cancel_event.is_set():
|
||
process_queue.put(None)
|
||
break
|
||
|
||
t0 = _time.perf_counter()
|
||
item = read_queue.get()
|
||
t_wait_total += _time.perf_counter() - t0
|
||
|
||
if item is None:
|
||
process_queue.put(None)
|
||
break
|
||
|
||
idx, frame = item
|
||
det_idx = idx - det_start_frame
|
||
frame_boxes = frames_detections[det_idx] if 0 <= det_idx < len(frames_detections) else []
|
||
|
||
if not frame_boxes:
|
||
process_queue.put((idx, frame))
|
||
frame_count += 1
|
||
continue
|
||
|
||
# 各人物ごとに個別ROIで処理（全員まとめると離れた人物間が巨大ROIになるため）
|
||
valid_boxes = []
|
||
for box in frame_boxes:
|
||
if not isinstance(box, list) or len(box) < 4:
|
||
continue
|
||
x, y, w, h = int(box[0]), int(box[1]), int(box[2]), int(box[3])
|
||
if w <= 0 or h <= 0:
|
||
continue
|
||
valid_boxes.append((x, y, w, h))
|
||
|
||
if not valid_boxes:
|
||
process_queue.put((idx, frame))
|
||
frame_count += 1
|
||
continue
|
||
|
||
for x, y, w, h in valid_boxes:
|
||
# display_scale で表示サイズを決定
|
||
cx = x + w / 2
|
||
cy = y + h / 2
|
||
dw = max(1, int(w * display_scale))
|
||
dh = max(1, int(h * display_scale))
|
||
dx = int(cx - dw / 2)
|
||
dy = int(cy - dh / 2)
|
||
|
||
# ROIは表示サイズ + blur_margin（計算用余白、境界問題回避のみ）
|
||
roi_x1 = max(0, dx - blur_margin)
|
||
roi_y1 = max(0, dy - blur_margin)
|
||
roi_x2 = min(src_width, dx + dw + blur_margin)
|
||
roi_y2 = min(src_height, dy + dh + blur_margin)
|
||
roi_width = roi_x2 - roi_x1
|
||
roi_height = roi_y2 - roi_y1
|
||
|
||
if roi_width <= 0 or roi_height <= 0:
|
||
continue
|
||
|
||
# ブラーはROI全体で計算（余白があるので端の精度が保証される）
|
||
roi_src = frame[roi_y1:roi_y2, roi_x1:roi_x2]
|
||
|
||
# ダウンサンプル→blur→アップサンプル（同等のぼかしを1/4の計算量で実現）
|
||
t1 = _time.perf_counter()
|
||
small_w = max(1, roi_width // 2)
|
||
small_h = max(1, roi_height // 2)
|
||
roi_small = cv2.resize(roi_src, (small_w, small_h), interpolation=cv2.INTER_LINEAR)
|
||
small_blur_size = max(3, (blur_size // 2) | 1)
|
||
roi_small_blurred = cv2.GaussianBlur(roi_small, (small_blur_size, small_blur_size), 0)
|
||
roi_blurred = cv2.resize(roi_small_blurred, (roi_width, roi_height), interpolation=cv2.INTER_LINEAR)
|
||
t_blur_total += _time.perf_counter() - t1
|
||
|
||
# 合成マスクはdisplay_scaleサイズの楕円のみ（featheringなし）
|
||
roi_mask = np.zeros((roi_height, roi_width), dtype=np.uint8)
|
||
center = (int(cx) - roi_x1, int(cy) - roi_y1)
|
||
axes = (max(1, dw // 2), max(1, dh // 2))
|
||
cv2.ellipse(roi_mask, center, axes, 0, 0, 360, 255, -1)
|
||
|
||
# バイナリマスクなのでcopyToで高速合成（float32変換不要）
|
||
t2 = _time.perf_counter()
|
||
result = roi_src.copy()
|
||
cv2.copyTo(roi_blurred, roi_mask, result)
|
||
frame[roi_y1:roi_y2, roi_x1:roi_x2] = result
|
||
t_blend_total += _time.perf_counter() - t2
|
||
|
||
process_queue.put((idx, frame))
|
||
frame_count += 1
|
||
|
||
if frame_count % REPORT_INTERVAL == 0:
|
||
n = max(frame_count, 1)
|
||
fps_proc = frame_count / max(t_wait_total + t_blur_total + t_blend_total, 1e-9)
|
||
print(
|
||
f"[Perf/Processor] frame={frame_count}"
|
||
f" wait={t_wait_total/n*1000:.1f}ms"
|
||
f" blur={t_blur_total/n*1000:.1f}ms"
|
||
f" blend={t_blend_total/n*1000:.1f}ms"
|
||
f" ROI={roi_width}x{roi_height}"
|
||
f" throughput≈{fps_proc:.1f}fps"
|
||
)
|
||
|
||
except Exception as e:
|
||
error_holder["error"] = f"Processor error: {e}"
|
||
process_queue.put(None)
|
||
finally:
|
||
if frame_count > 0:
|
||
n = max(frame_count, 1)
|
||
print(
|
||
f"[Perf/Processor] FINAL frame={frame_count}"
|
||
f" wait_avg={t_wait_total/n*1000:.1f}ms"
|
||
f" blur_avg={t_blur_total/n*1000:.1f}ms"
|
||
f" blend_avg={t_blend_total/n*1000:.1f}ms"
|
||
)
|
||
|
||
def _writer_worker():
|
||
"""Write frames to output."""
|
||
import time as _time
|
||
t_wait_total = 0.0
|
||
t_write_total = 0.0
|
||
frame_count = 0
|
||
writer = None
|
||
try:
|
||
writer = _build_video_writer(req.output_path, req.format, src_fps, src_width, src_height)
|
||
|
||
while True:
|
||
if cancel_event and cancel_event.is_set():
|
||
break
|
||
|
||
t0 = _time.perf_counter()
|
||
item = process_queue.get()
|
||
t_wait_total += _time.perf_counter() - t0
|
||
|
||
if item is None:
|
||
break
|
||
|
||
idx, frame = item
|
||
t1 = _time.perf_counter()
|
||
writer.write(frame)
|
||
t_write_total += _time.perf_counter() - t1
|
||
frame_count += 1
|
||
|
||
with progress_lock:
|
||
current_progress[0] = idx + 1
|
||
tasks[task_id].progress = current_progress[0]
|
||
|
||
except Exception as e:
|
||
error_holder["error"] = f"Writer error: {e}"
|
||
finally:
|
||
if writer:
|
||
try:
|
||
writer.release()
|
||
except Exception as e:
|
||
print(f"[FaceMask] Writer release error: {e}")
|
||
if frame_count > 0:
|
||
n = max(frame_count, 1)
|
||
print(
|
||
f"[Perf/Writer] FINAL frame={frame_count}"
|
||
f" wait_avg={t_wait_total/n*1000:.1f}ms"
|
||
f" write_avg={t_write_total/n*1000:.1f}ms"
|
||
)
|
||
|
||
print(
|
||
f"[FaceMask] Starting blur bake: {req.video_path} + "
|
||
f"{req.detections_path} -> {req.output_path}"
|
||
)
|
||
|
||
# Start threads
|
||
reader_thread = threading.Thread(target=_reader_worker, daemon=True)
|
||
processor_thread = threading.Thread(target=_processor_worker, daemon=True)
|
||
writer_thread = threading.Thread(target=_writer_worker, daemon=True)
|
||
|
||
reader_thread.start()
|
||
processor_thread.start()
|
||
writer_thread.start()
|
||
|
||
# Wait for completion
|
||
reader_thread.join()
|
||
processor_thread.join()
|
||
writer_thread.join()
|
||
|
||
if error_holder["error"]:
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = error_holder["error"]
|
||
print(f"[FaceMask] Bake failed: {error_holder['error']}")
|
||
elif cancel_event and cancel_event.is_set():
|
||
tasks[task_id].status = TaskStatus.CANCELLED
|
||
tasks[task_id].message = "Cancelled by user"
|
||
else:
|
||
tasks[task_id].status = TaskStatus.COMPLETED
|
||
tasks[task_id].result_path = req.output_path
|
||
tasks[task_id].message = "Blur bake completed"
|
||
print(f"[FaceMask] Bake completed: {req.output_path}")
|
||
|
||
except Exception as e:
|
||
tasks[task_id].status = TaskStatus.FAILED
|
||
tasks[task_id].message = str(e)
|
||
print(f"Error in async bake task {task_id}: {e}")
|
||
traceback.print_exc()
|
||
finally:
|
||
if task_id in cancel_events:
|
||
del cancel_events[task_id]
|
||
|
||
|
||
def check_gpu_available() -> dict:
|
||
"""
|
||
Check if GPU is available for inference.
|
||
|
||
Returns a dict with GPU information:
|
||
{
|
||
"available": bool,
|
||
"device_name": str or None,
|
||
"device_count": int,
|
||
"rocm_version": str or None
|
||
}
|
||
"""
|
||
global _gpu_status_cache
|
||
|
||
# Return cached result if available
|
||
if _gpu_status_cache is not None:
|
||
return _gpu_status_cache
|
||
|
||
result = {
|
||
"available": False,
|
||
"device_name": None,
|
||
"device_count": 0,
|
||
"rocm_version": None
|
||
}
|
||
|
||
try:
|
||
import torch
|
||
|
||
result["available"] = torch.cuda.is_available()
|
||
|
||
if result["available"]:
|
||
result["device_count"] = torch.cuda.device_count()
|
||
if result["device_count"] > 0:
|
||
result["device_name"] = torch.cuda.get_device_name(0)
|
||
|
||
if hasattr(torch.version, 'hip'):
|
||
result["rocm_version"] = torch.version.hip
|
||
|
||
except Exception as e:
|
||
print(f"[FaceMask] Warning: GPU detection failed: {e}")
|
||
result["available"] = False
|
||
|
||
# Cache the result
|
||
_gpu_status_cache = result
|
||
|
||
return result
|
||
|
||
|
||
def log_startup_diagnostics():
|
||
"""Log diagnostic information about the environment and GPU."""
|
||
print("=" * 70)
|
||
print("[FaceMask Server] Startup Diagnostics")
|
||
print("=" * 70)
|
||
|
||
# Python Environment
|
||
print("\n[Python Environment]")
|
||
print(f" Python Version: {sys.version.split()[0]}")
|
||
print(f" Python Executable: {sys.executable}")
|
||
print(f" Platform: {platform.platform()}")
|
||
print(f" Working Directory: {os.getcwd()}")
|
||
|
||
# Check if in venv
|
||
in_venv = sys.prefix != sys.base_prefix
|
||
print(f" Virtual Environment: {'Yes' if in_venv else 'No'}")
|
||
if in_venv:
|
||
print(f" venv path: {sys.prefix}")
|
||
|
||
# ROCm Environment Variables
|
||
print("\n[ROCm Environment Variables]")
|
||
rocm_vars = [
|
||
'ROCM_PATH',
|
||
'HSA_OVERRIDE_GFX_VERSION',
|
||
'PYTORCH_ROCM_ARCH',
|
||
'ROCBLAS_TENSILE_LIBPATH',
|
||
'LD_LIBRARY_PATH'
|
||
]
|
||
|
||
for var in rocm_vars:
|
||
value = os.environ.get(var)
|
||
if value:
|
||
# 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(" ⚠️ WARNING: ROCm library paths not found!")
|
||
else:
|
||
if len(value) > 200:
|
||
display_value = value[:200] + "... (truncated)"
|
||
else:
|
||
display_value = value
|
||
print(f" {var}: {display_value}")
|
||
else:
|
||
print(f" {var}: (not set)")
|
||
|
||
# GPU Detection
|
||
print("\n[GPU Detection]")
|
||
try:
|
||
import torch
|
||
|
||
cuda_available = torch.cuda.is_available()
|
||
print(f" torch.cuda.is_available(): {cuda_available}")
|
||
|
||
if cuda_available:
|
||
device_count = torch.cuda.device_count()
|
||
print(f" GPU Device Count: {device_count}")
|
||
|
||
if device_count > 0:
|
||
device_name = torch.cuda.get_device_name(0)
|
||
print(f" GPU Device 0: {device_name}")
|
||
|
||
# ROCm version
|
||
if hasattr(torch.version, 'hip'):
|
||
print(f" ROCm Version (HIP): {torch.version.hip}")
|
||
|
||
# CUDA version (might be emulated by ROCm)
|
||
if torch.version.cuda:
|
||
print(f" CUDA Version: {torch.version.cuda}")
|
||
else:
|
||
print(" WARNING: GPU not detected!")
|
||
print(" Server will use CPU for inference (slower)")
|
||
print(" Troubleshooting:")
|
||
print(" - Check ROCm environment variables above")
|
||
print(" - Run: python -c 'import torch; print(torch.cuda.is_available())'")
|
||
|
||
except ImportError as e:
|
||
print(f" ERROR: Cannot import torch: {e}")
|
||
print(" PyTorch must be installed for inference")
|
||
except Exception as e:
|
||
print(f" ERROR during GPU detection: {e}")
|
||
|
||
print("=" * 70)
|
||
print()
|
||
|
||
|
||
@app.get("/status")
|
||
def get_status():
|
||
gpu_info = check_gpu_available()
|
||
|
||
return {
|
||
"status": "running",
|
||
"gpu_available": gpu_info["available"],
|
||
"gpu_device": gpu_info["device_name"],
|
||
"gpu_count": gpu_info["device_count"],
|
||
"rocm_version": gpu_info["rocm_version"]
|
||
}
|
||
|
||
|
||
@app.post("/video_info")
|
||
def get_video_info(req: VideoInfoRequest):
|
||
if not os.path.exists(req.video_path):
|
||
raise HTTPException(status_code=404, detail=f"Video not found: {req.video_path}")
|
||
|
||
cap = cv2.VideoCapture(req.video_path)
|
||
if not cap.isOpened():
|
||
raise HTTPException(status_code=400, detail="Failed to open video")
|
||
|
||
try:
|
||
fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0)
|
||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
|
||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
|
||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
|
||
finally:
|
||
cap.release()
|
||
|
||
return {
|
||
"video_path": req.video_path,
|
||
"fps": fps,
|
||
"width": width,
|
||
"height": height,
|
||
"frame_count": frame_count,
|
||
}
|
||
|
||
@app.post("/generate", response_model=Task)
|
||
def generate_mask_endpoint(req: GenerateRequest, background_tasks: BackgroundTasks):
|
||
task_id = str(uuid.uuid4())
|
||
task = Task(id=task_id, status=TaskStatus.PENDING)
|
||
tasks[task_id] = task
|
||
cancel_events[task_id] = threading.Event()
|
||
|
||
background_tasks.add_task(process_video_task, task_id, req)
|
||
return task
|
||
|
||
|
||
@app.post("/bake_blur", response_model=Task)
|
||
def bake_blur_endpoint(req: BakeRequest, background_tasks: BackgroundTasks):
|
||
task_id = str(uuid.uuid4())
|
||
task = Task(id=task_id, status=TaskStatus.PENDING)
|
||
tasks[task_id] = task
|
||
cancel_events[task_id] = threading.Event()
|
||
|
||
background_tasks.add_task(process_bake_task, task_id, req)
|
||
return task
|
||
|
||
@app.get("/tasks/{task_id}", response_model=Task)
|
||
def get_task(task_id: str):
|
||
if task_id not in tasks:
|
||
raise HTTPException(status_code=404, detail="Task not found")
|
||
return tasks[task_id]
|
||
|
||
@app.post("/tasks/{task_id}/cancel")
|
||
def cancel_task(task_id: str):
|
||
if task_id not in tasks:
|
||
raise HTTPException(status_code=404, detail="Task not found")
|
||
|
||
if task_id in cancel_events:
|
||
cancel_events[task_id].set()
|
||
|
||
return {"message": "Cancellation requested"}
|
||
|
||
if __name__ == "__main__":
|
||
log_startup_diagnostics()
|
||
uvicorn.run(app, host="127.0.0.1", port=8181)
|