20260430_001744-feat: PyTorch CUDA + SAM2 GPU inference, video thumbnail, real FPS + configurable parse FPS, DICOM batch import
This commit is contained in:
@@ -5,7 +5,7 @@ import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
@@ -16,12 +16,43 @@ from minio_client import upload_file, BUCKET_NAME
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_video_fps(video_path: str) -> float:
|
||||
"""Read the original frame rate of a video file."""
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
return 30.0
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
cap.release()
|
||||
return fps if fps > 0 else 30.0
|
||||
|
||||
|
||||
def extract_thumbnail(video_path: str, output_path: str, width: int = 640) -> str:
|
||||
"""Extract the first frame of a video as a thumbnail JPEG."""
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
raise RuntimeError(f"Cannot open video for thumbnail: {video_path}")
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
if not ret or frame is None:
|
||||
raise RuntimeError(f"Cannot read first frame from: {video_path}")
|
||||
|
||||
h, w = frame.shape[:2]
|
||||
if w > width:
|
||||
scale = width / w
|
||||
new_w = int(w * scale)
|
||||
new_h = int(h * scale)
|
||||
frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
cv2.imwrite(output_path, frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
|
||||
return output_path
|
||||
|
||||
|
||||
def parse_video(
|
||||
video_path: str,
|
||||
output_dir: str,
|
||||
fps: int = 30,
|
||||
max_frames: Optional[int] = None,
|
||||
) -> List[str]:
|
||||
) -> Tuple[List[str], float]:
|
||||
"""Extract frames from a video file using FFmpeg or OpenCV fallback.
|
||||
|
||||
Args:
|
||||
@@ -31,10 +62,11 @@ def parse_video(
|
||||
max_frames: Optional maximum number of frames to extract.
|
||||
|
||||
Returns:
|
||||
List of paths to extracted frame images.
|
||||
Tuple of (frame_paths, original_fps).
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
frame_paths: List[str] = []
|
||||
original_fps = get_video_fps(video_path)
|
||||
|
||||
# Try FFmpeg first
|
||||
if shutil.which("ffmpeg"):
|
||||
@@ -57,7 +89,7 @@ def parse_video(
|
||||
if max_frames:
|
||||
frame_paths = frame_paths[:max_frames]
|
||||
logger.info("Extracted %d frames via FFmpeg", len(frame_paths))
|
||||
return frame_paths
|
||||
return frame_paths, original_fps
|
||||
else:
|
||||
logger.warning("FFmpeg failed: %s", result.stderr)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
@@ -89,7 +121,7 @@ def parse_video(
|
||||
|
||||
cap.release()
|
||||
logger.info("Extracted %d frames via OpenCV", len(frame_paths))
|
||||
return frame_paths
|
||||
return frame_paths, original_fps
|
||||
|
||||
|
||||
def parse_dicom(
|
||||
@@ -134,12 +166,12 @@ def parse_dicom(
|
||||
# Handle multi-frame DICOM
|
||||
if pixel_array.ndim == 3:
|
||||
for f in range(pixel_array.shape[0]):
|
||||
out_path = os.path.join(output_dir, f"frame_{idx:06d}_{f:03d}.png")
|
||||
cv2.imwrite(out_path, pixel_array[f])
|
||||
out_path = os.path.join(output_dir, f"frame_{idx:06d}_{f:03d}.jpg")
|
||||
cv2.imwrite(out_path, pixel_array[f], [cv2.IMWRITE_JPEG_QUALITY, 85])
|
||||
frame_paths.append(out_path)
|
||||
else:
|
||||
out_path = os.path.join(output_dir, f"frame_{idx:06d}.png")
|
||||
cv2.imwrite(out_path, pixel_array)
|
||||
out_path = os.path.join(output_dir, f"frame_{idx:06d}.jpg")
|
||||
cv2.imwrite(out_path, pixel_array, [cv2.IMWRITE_JPEG_QUALITY, 85])
|
||||
frame_paths.append(out_path)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error("Failed to read DICOM %s: %s", path, exc)
|
||||
|
||||
Reference in New Issue
Block a user