290 lines
9.9 KiB
Python
290 lines
9.9 KiB
Python
"""Media upload and parsing endpoints."""
|
|
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import List, Optional
|
|
|
|
from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile, status
|
|
from sqlalchemy.orm import Session
|
|
|
|
from database import get_db
|
|
from minio_client import upload_file, get_presigned_url, download_file
|
|
from models import Project, Frame
|
|
from schemas import FrameOut
|
|
from services.frame_parser import (
|
|
parse_video, parse_dicom, upload_frames_to_minio,
|
|
extract_thumbnail, get_video_fps,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
router = APIRouter(prefix="/api/media", tags=["Media"])
|
|
|
|
ALLOWED_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".png", ".jpg", ".jpeg", ".dcm"}
|
|
|
|
|
|
def _get_ext(filename: str) -> str:
|
|
return Path(filename).suffix.lower()
|
|
|
|
|
|
@router.post(
|
|
"/upload",
|
|
status_code=status.HTTP_201_CREATED,
|
|
summary="Upload a media file",
|
|
)
|
|
async def upload_media(
|
|
file: UploadFile = File(...),
|
|
project_id: Optional[int] = Form(None),
|
|
db: Session = Depends(get_db),
|
|
) -> dict:
|
|
"""Accept a video, image, or DICOM file and store it in MinIO.
|
|
|
|
If project_id is provided, the video_path of the project is updated.
|
|
Returns the presigned URL of the uploaded object.
|
|
"""
|
|
if not file.filename:
|
|
raise HTTPException(status_code=400, detail="Missing filename")
|
|
|
|
ext = _get_ext(file.filename)
|
|
if ext not in ALLOWED_EXTENSIONS:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Unsupported file type: {ext}",
|
|
)
|
|
|
|
data = await file.read()
|
|
object_name = f"uploads/{project_id or 'general'}/{file.filename}"
|
|
|
|
try:
|
|
upload_file(object_name, data, content_type=file.content_type or "application/octet-stream", length=len(data))
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.error("Upload failed: %s", exc)
|
|
raise HTTPException(status_code=500, detail="Upload to storage failed") from exc
|
|
|
|
file_url = get_presigned_url(object_name, expires=3600)
|
|
|
|
if project_id:
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if project:
|
|
project.video_path = object_name
|
|
db.commit()
|
|
logger.info("Linked upload to project_id=%s", project_id)
|
|
else:
|
|
logger.warning("Project id=%s not found for upload linkage", project_id)
|
|
else:
|
|
# Auto-create a project named after the file
|
|
project = Project(
|
|
name=file.filename,
|
|
description="Auto-created from upload",
|
|
status="pending",
|
|
video_path=object_name,
|
|
source_type="video",
|
|
)
|
|
db.add(project)
|
|
db.commit()
|
|
db.refresh(project)
|
|
project_id = project.id
|
|
object_name = f"uploads/{project_id}/{file.filename}"
|
|
# Re-upload with corrected path
|
|
upload_file(object_name, data, content_type=file.content_type or "application/octet-stream", length=len(data))
|
|
project.video_path = object_name
|
|
db.commit()
|
|
logger.info("Auto-created project id=%s for upload %s", project_id, file.filename)
|
|
|
|
logger.info("Upload complete: %s (size=%d bytes). Async parsing queued.", object_name, len(data))
|
|
|
|
return {
|
|
"object_name": object_name,
|
|
"file_url": file_url,
|
|
"size": len(data),
|
|
"project_id": project_id,
|
|
"message": "Upload successful. Parsing job queued.",
|
|
}
|
|
|
|
|
|
@router.post(
|
|
"/upload/dicom",
|
|
status_code=status.HTTP_201_CREATED,
|
|
summary="Upload multiple DICOM files",
|
|
)
|
|
async def upload_dicom_batch(
|
|
files: List[UploadFile] = File(...),
|
|
project_id: Optional[int] = Form(None),
|
|
db: Session = Depends(get_db),
|
|
) -> dict:
|
|
"""Upload multiple .dcm files for a DICOM series.
|
|
|
|
If project_id is provided, files are added to the existing project.
|
|
Otherwise a new DICOM project is created.
|
|
"""
|
|
if not files:
|
|
raise HTTPException(status_code=400, detail="No files uploaded")
|
|
|
|
uploaded = []
|
|
|
|
if project_id:
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
else:
|
|
# Create new DICOM project
|
|
first_name = files[0].filename or "DICOM_Series"
|
|
project = Project(
|
|
name=first_name,
|
|
description=f"DICOM series with {len(files)} files",
|
|
status="pending",
|
|
source_type="dicom",
|
|
)
|
|
db.add(project)
|
|
db.commit()
|
|
db.refresh(project)
|
|
project_id = project.id
|
|
logger.info("Auto-created DICOM project id=%s", project_id)
|
|
|
|
for file in files:
|
|
if not file.filename or not file.filename.lower().endswith(".dcm"):
|
|
continue
|
|
data = await file.read()
|
|
object_name = f"uploads/{project_id}/dicom/{file.filename}"
|
|
try:
|
|
upload_file(object_name, data, content_type="application/dicom", length=len(data))
|
|
uploaded.append(object_name)
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.error("Failed to upload DICOM %s: %s", file.filename, exc)
|
|
|
|
project.video_path = f"uploads/{project_id}/dicom"
|
|
db.commit()
|
|
|
|
return {
|
|
"project_id": project_id,
|
|
"uploaded_count": len(uploaded),
|
|
"message": f"Uploaded {len(uploaded)} DICOM files. Parsing job queued.",
|
|
}
|
|
|
|
|
|
@router.post(
|
|
"/parse",
|
|
status_code=status.HTTP_202_ACCEPTED,
|
|
summary="Trigger frame extraction",
|
|
)
|
|
def parse_media(
|
|
project_id: int,
|
|
source_type: Optional[str] = None,
|
|
db: Session = Depends(get_db),
|
|
) -> dict:
|
|
"""Trigger frame extraction for a project's uploaded media.
|
|
|
|
* video: uses FFmpeg or OpenCV fallback, extracts thumbnail.
|
|
* dicom: uses pydicom to read DCM frames.
|
|
|
|
Extracted frames are uploaded to MinIO and registered in the database.
|
|
"""
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
if not project.video_path:
|
|
raise HTTPException(status_code=400, detail="Project has no media uploaded")
|
|
|
|
effective_source = source_type or project.source_type or "video"
|
|
parse_fps = project.parse_fps or 30.0
|
|
|
|
tmp_dir = tempfile.mkdtemp(prefix=f"seg_parse_{project_id}_")
|
|
output_dir = os.path.join(tmp_dir, "frames")
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
try:
|
|
if effective_source == "dicom":
|
|
# Download all dicom files from MinIO
|
|
dcm_dir = os.path.join(tmp_dir, "dcm")
|
|
os.makedirs(dcm_dir, exist_ok=True)
|
|
|
|
from minio_client import get_minio_client, BUCKET_NAME
|
|
client = get_minio_client()
|
|
prefix = project.video_path
|
|
objects = list(client.list_objects(BUCKET_NAME, prefix=prefix, recursive=True))
|
|
for obj in objects:
|
|
if obj.object_name.lower().endswith(".dcm"):
|
|
data = download_file(obj.object_name)
|
|
local_dcm = os.path.join(dcm_dir, os.path.basename(obj.object_name))
|
|
with open(local_dcm, "wb") as f:
|
|
f.write(data)
|
|
|
|
frame_files = parse_dicom(dcm_dir, output_dir)
|
|
else:
|
|
# Video: download and parse
|
|
media_bytes = download_file(project.video_path)
|
|
local_path = os.path.join(tmp_dir, Path(project.video_path).name)
|
|
with open(local_path, "wb") as f:
|
|
f.write(media_bytes)
|
|
|
|
frame_files, original_fps = parse_video(local_path, output_dir, fps=int(parse_fps))
|
|
project.original_fps = original_fps
|
|
|
|
# Extract thumbnail from first frame
|
|
thumbnail_path = os.path.join(tmp_dir, "thumbnail.jpg")
|
|
try:
|
|
extract_thumbnail(local_path, thumbnail_path)
|
|
with open(thumbnail_path, "rb") as f:
|
|
thumb_data = f.read()
|
|
thumb_object = f"projects/{project_id}/thumbnail.jpg"
|
|
upload_file(thumb_object, thumb_data, content_type="image/jpeg", length=len(thumb_data))
|
|
project.thumbnail_url = thumb_object
|
|
logger.info("Uploaded thumbnail for project_id=%s", project_id)
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.warning("Thumbnail extraction failed: %s", exc)
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.error("Frame extraction failed: %s", exc)
|
|
shutil.rmtree(tmp_dir, ignore_errors=True)
|
|
raise HTTPException(status_code=500, detail="Frame extraction failed") from exc
|
|
|
|
# Upload frames to MinIO
|
|
try:
|
|
object_names = upload_frames_to_minio(frame_files, project_id)
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.error("Frame upload failed: %s", exc)
|
|
shutil.rmtree(tmp_dir, ignore_errors=True)
|
|
raise HTTPException(status_code=500, detail="Frame upload to storage failed") from exc
|
|
|
|
# Register frames in DB
|
|
frames_out = []
|
|
for idx, obj_name in enumerate(object_names):
|
|
local_frame = frame_files[idx]
|
|
try:
|
|
import cv2
|
|
img = cv2.imread(local_frame)
|
|
h, w = img.shape[:2] if img is not None else (None, None)
|
|
except Exception: # noqa: BLE001
|
|
h, w = None, None
|
|
|
|
frame = Frame(
|
|
project_id=project_id,
|
|
frame_index=idx,
|
|
image_url=obj_name,
|
|
width=w,
|
|
height=h,
|
|
)
|
|
db.add(frame)
|
|
frames_out.append(frame)
|
|
|
|
db.commit()
|
|
for f in frames_out:
|
|
db.refresh(f)
|
|
|
|
# Cleanup temp files
|
|
shutil.rmtree(tmp_dir, ignore_errors=True)
|
|
|
|
project.status = "ready"
|
|
db.commit()
|
|
|
|
logger.info("Parsed %d frames for project_id=%s", len(frames_out), project_id)
|
|
return {
|
|
"project_id": project_id,
|
|
"frames_extracted": len(frames_out),
|
|
"status": "ready",
|
|
"message": "Frame extraction completed successfully.",
|
|
}
|