update the codebase poc ver1
This commit is contained in:
DatTT127
@@ -0,0 +1,48 @@
|
||||
import logging
|
||||
from typing import NamedTuple, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class DriftResult:
|
||||
score: float
|
||||
is_drifted: bool
|
||||
threshold: float
|
||||
|
||||
@dataclass
|
||||
class GuardrailResult:
|
||||
verdict: str # "PASS" | "MITIGATE"
|
||||
reason: Optional[str] = None
|
||||
|
||||
class BERTAdapter:
|
||||
"""
|
||||
Adapter for BERT-based safety checks (drift, referee, guardrails).
|
||||
Current implementation provides stubs.
|
||||
"""
|
||||
def __init__(self):
|
||||
self.drift_threshold = 0.7
|
||||
|
||||
def drift_check(self, text: str) -> DriftResult:
|
||||
"""
|
||||
Checks if the input text drifts from the clinical domain.
|
||||
"""
|
||||
# Stub: Always return no drift
|
||||
return DriftResult(score=0.9, is_drifted=False, threshold=self.drift_threshold)
|
||||
|
||||
def referee_check(self, text: str, retrieved_chunks: list) -> bool:
|
||||
"""
|
||||
RAG-Referee: Validates if LLM response is grounded in provided chunks.
|
||||
"""
|
||||
# Stub: Always return grounded
|
||||
return True
|
||||
|
||||
def guardrail_check(self, text: str) -> GuardrailResult:
|
||||
"""
|
||||
Token/Chunk level guardrail check for hallucinations or scope violations.
|
||||
"""
|
||||
# Stub: Always return PASS
|
||||
return GuardrailResult(verdict="PASS")
|
||||
|
||||
def get_bert_adapter() -> BERTAdapter:
|
||||
return BERTAdapter()
|
||||
@@ -0,0 +1,71 @@
|
||||
import logging
|
||||
from typing import Any, AsyncGenerator, List, Optional
|
||||
from langchain_google_vertexai import VertexAI
|
||||
from langchain.callbacks.base import BaseCallbackHandler
|
||||
from langchain.schema import LLMResult
|
||||
from backend.implementation import config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class AuditCallbackHandler(BaseCallbackHandler):
|
||||
"""
|
||||
Langchain callback to enforce audit logging before LLM calls per NFR-16a.
|
||||
"""
|
||||
def __init__(self, session_id: str, metadata: Optional[dict] = None):
|
||||
self.session_id = session_id
|
||||
self.metadata = metadata or {}
|
||||
|
||||
def on_llm_start(self, serialized: Any, prompts: List[str], **kwargs) -> None:
|
||||
# MANDATORY: Write egress_consent + egress_redact_manifest to immutable audit log
|
||||
# In a real implementation, this would call a database service to commit to Postgres.
|
||||
logger.info(f"[AUDIT] Pre-egress audit commit for session {self.session_id}. "
|
||||
f"Prompts: {prompts}. Metadata: {self.metadata}")
|
||||
|
||||
def on_llm_end(self, response: LLMResult, **kwargs) -> None:
|
||||
# Log actual egress event after completion
|
||||
logger.info(f"[AUDIT] LLM egress completed for session {self.session_id}")
|
||||
|
||||
def on_llm_error(self, error: Exception, **kwargs) -> None:
|
||||
logger.error(f"[AUDIT] LLM error for session {self.session_id}: {str(error)}")
|
||||
|
||||
class VertexAILangchainAdapter:
|
||||
def __init__(self):
|
||||
self.llm = VertexAI(
|
||||
model_name=config.VERTEX_AI_MODEL,
|
||||
project_id=config.VERTEX_AI_PROJECT,
|
||||
location=config.VERTEX_AI_LOCATION,
|
||||
max_output_tokens=256,
|
||||
temperature=0.2,
|
||||
top_p=0.8,
|
||||
top_k=40,
|
||||
)
|
||||
|
||||
async def generate(self, prompt: str, session_id: str, metadata: Optional[dict] = None) -> str:
|
||||
import asyncio
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
callback_handler = AuditCallbackHandler(session_id, metadata)
|
||||
|
||||
def _sync_generate():
|
||||
result = self.llm.generate(
|
||||
prompts=[prompt],
|
||||
callbacks=[callback_handler]
|
||||
)
|
||||
return result.generations[0][0].text
|
||||
|
||||
return await loop.run_in_executor(None, _sync_generate)
|
||||
|
||||
async def stream_generate(self, prompt: str, session_id: str, metadata: Optional[dict] = None) -> AsyncGenerator[str, None]:
|
||||
import asyncio
|
||||
loop = asyncio.get_event_loop()
|
||||
callback_handler = AuditCallbackHandler(session_id, metadata)
|
||||
|
||||
def _sync_stream():
|
||||
return self.llm.stream(prompt, callbacks=[callback_handler])
|
||||
|
||||
stream = await loop.run_in_executor(None, _sync_stream)
|
||||
for chunk in stream:
|
||||
yield chunk
|
||||
|
||||
def get_llm_adapter() -> VertexAILangchainAdapter:
|
||||
return VertexAILangchainAdapter()
|
||||
@@ -0,0 +1,40 @@
|
||||
import redis
|
||||
import logging
|
||||
from backend.implementation import config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class RedisClient:
|
||||
"""
|
||||
Singleton Redis client for managing session state and consult_mode.
|
||||
"""
|
||||
_instance = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super(RedisClient, cls).__new__(cls)
|
||||
try:
|
||||
cls._instance.client = redis.Redis(
|
||||
host=config.REDIS_HOST,
|
||||
port=config.REDIS_PORT,
|
||||
db=config.REDIS_DB,
|
||||
decode_responses=True
|
||||
)
|
||||
logger.info("Connected to Redis at %s:%s", config.REDIS_HOST, config.REDIS_PORT)
|
||||
except Exception as e:
|
||||
logger.error("Failed to connect to Redis: %s", e)
|
||||
cls._instance.client = None
|
||||
return cls._instance
|
||||
|
||||
def get(self, key: str):
|
||||
return self.client.get(key) if self.client else None
|
||||
|
||||
def set(self, key: str, value: str, ex: int = None):
|
||||
if self.client:
|
||||
self.client.set(key, value, ex=ex)
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
return bool(self.client.exists(key)) if self.client else False
|
||||
|
||||
def get_redis_client() -> RedisClient:
|
||||
return RedisClient()
|
||||
@@ -0,0 +1,121 @@
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
import numpy as np
|
||||
import requests
|
||||
|
||||
|
||||
class TritonAdapter:
|
||||
def __init__(self, endpoint_url: str, timeout: float = 60.0):
|
||||
self.endpoint_url = endpoint_url.rstrip("/")
|
||||
self.timeout = timeout
|
||||
|
||||
|
||||
async def close(self):
|
||||
pass
|
||||
|
||||
|
||||
async def infer(
|
||||
self, model_name: str, inputs: dict, outputs: list[str] | None = None
|
||||
) -> dict:
|
||||
return await asyncio.to_thread(
|
||||
self._infer_sync, model_name, inputs, outputs
|
||||
)
|
||||
|
||||
def _infer_sync(
|
||||
self, model_name: str, inputs: dict, outputs: list[str] | None = None
|
||||
) -> dict:
|
||||
metadata_inputs = []
|
||||
binary_chunks = []
|
||||
|
||||
for name, spec in inputs.items():
|
||||
data = spec["data"]
|
||||
shape = spec.get("shape", [])
|
||||
datatype = spec.get("datatype", "FP32")
|
||||
|
||||
arr = np.asarray(data, dtype=np.float32)
|
||||
binary = arr.tobytes()
|
||||
|
||||
metadata_inputs.append(
|
||||
{
|
||||
"name": name,
|
||||
"shape": shape,
|
||||
"datatype": datatype,
|
||||
"parameters": {"binary_data_size": len(binary)},
|
||||
}
|
||||
)
|
||||
binary_chunks.append(binary)
|
||||
|
||||
metadata_outputs = [{"name": o} for o in (outputs or [])]
|
||||
metadata = {
|
||||
"inputs": metadata_inputs,
|
||||
"outputs": metadata_outputs,
|
||||
}
|
||||
|
||||
metadata_bytes = json.dumps(metadata).encode("utf-8")
|
||||
body = metadata_bytes + b"".join(binary_chunks)
|
||||
|
||||
headers = {
|
||||
"Inference-Header-Content-Length": str(len(metadata_bytes)),
|
||||
"Content-Type": "application/octet-stream",
|
||||
}
|
||||
|
||||
url = f"{self.endpoint_url}/v2/models/{model_name}/infer"
|
||||
response = requests.post(url, data=body, headers=headers, timeout=self.timeout)
|
||||
response.raise_for_status()
|
||||
return self._parse_binary_response(response.headers, response.content)
|
||||
|
||||
@staticmethod
|
||||
def _parse_binary_response(headers: dict, body: bytes) -> dict:
|
||||
header_len = int(headers.get("Inference-Header-Content-Length", "0"))
|
||||
metadata = json.loads(body[:header_len].decode("utf-8"))
|
||||
|
||||
result = {}
|
||||
offset = 0
|
||||
for output in metadata.get("outputs", []):
|
||||
name = output["name"]
|
||||
shape = output.get("shape", [])
|
||||
params = output.get("parameters", {})
|
||||
binary_size = params.get("binary_data_size", 0)
|
||||
|
||||
if binary_size > 0:
|
||||
chunk = body[header_len + offset : header_len + offset + binary_size]
|
||||
arr = np.frombuffer(chunk, dtype=np.float32).reshape(shape)
|
||||
result[name] = arr.tolist()
|
||||
offset += binary_size
|
||||
|
||||
return result
|
||||
|
||||
async def model_ready(self, model_name: str) -> bool:
|
||||
return await asyncio.to_thread(self._model_ready_sync, model_name)
|
||||
|
||||
def _model_ready_sync(self, model_name: str) -> bool:
|
||||
url = f"{self.endpoint_url}/v2/models/{model_name}"
|
||||
response = requests.get(url, timeout=self.timeout)
|
||||
if response.status_code == 404:
|
||||
return False
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data.get("ready", False)
|
||||
|
||||
async def list_models(self) -> list[dict]:
|
||||
return await asyncio.to_thread(self._list_models_sync)
|
||||
|
||||
# def _list_models_sync(self) -> list[dict]:
|
||||
# url = f"{self.endpoint_url}/v2/models"
|
||||
# response = requests.get(url, timeout=self.timeout)
|
||||
# response.raise_for_status()
|
||||
# data = response.json()
|
||||
# return data.get("models", [])
|
||||
|
||||
def _list_models_sync(self) -> list[dict]:
|
||||
# 1. Change the endpoint to Triton's repository index path
|
||||
url = f"{self.endpoint_url}/v2/repository/index"
|
||||
|
||||
# 2. Change requests.get to requests.post with an empty json payload {}
|
||||
response = requests.post(url, json={}, timeout=self.timeout)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
# KServe v2 returns a list directly: [{"name": "model_a", "version": "1", "state": "READY"}]
|
||||
return data
|
||||
@@ -0,0 +1,323 @@
|
||||
import asyncio
|
||||
import io
|
||||
import base64
|
||||
import uuid
|
||||
import logging
|
||||
import numpy as np
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from data.spec.schemas import (
|
||||
AnalysisJobSubmit, JobStatus, JobResult, PipelineStep,
|
||||
StepEvent, ModelCatalog, ModelRegistrationResult,
|
||||
HealthStatus,
|
||||
)
|
||||
from PIL import Image
|
||||
|
||||
from backend.implementation.preprocessing.clahe import apply_clahe
|
||||
from backend.implementation.preprocessing.tensor_prep import (
|
||||
prepare_angle_tensor,
|
||||
prepare_inflammation_tensor,
|
||||
prepare_segmentation_tensor,
|
||||
)
|
||||
from backend.implementation.postprocessing.measurement import calculate_thickness
|
||||
from backend.implementation.postprocessing.severity import calculate_severity
|
||||
from backend.implementation.postprocessing.overlay import create_overlay
|
||||
from backend.implementation.postprocessing.calibration import (
|
||||
calibration_config_from_params,
|
||||
interpret_angle_logits,
|
||||
interpret_inflammation_logits,
|
||||
)
|
||||
from backend.implementation.config import (
|
||||
get_model_name,
|
||||
get_segmentation_model,
|
||||
get_angle_type,
|
||||
TRITON_ENDPOINT,
|
||||
)
|
||||
from backend.implementation.adapters.triton_adapter import TritonAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_job_registry: dict[str, dict] = {}
|
||||
_job_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _interpret_angle_result(result: dict, params: dict | None = None) -> dict:
|
||||
logits = result.get("logits", [])
|
||||
if not logits:
|
||||
raise ValueError("Empty angle logits")
|
||||
config = calibration_config_from_params(params)
|
||||
return interpret_angle_logits(logits, config)
|
||||
|
||||
|
||||
def _interpret_inflammation_result(result: dict, params: dict | None = None) -> dict:
|
||||
logits = result.get("logits", [])
|
||||
if not logits:
|
||||
raise ValueError("Empty inflammation logits")
|
||||
config = calibration_config_from_params(params)
|
||||
return interpret_inflammation_logits(logits, config)
|
||||
|
||||
|
||||
def _process_segmentation_result(result: dict, angle_class: str) -> tuple:
|
||||
logits = result.get("logits", [])
|
||||
if not logits:
|
||||
raise ValueError("Empty segmentation logits")
|
||||
logits_arr = np.array(logits)
|
||||
if logits_arr.ndim < 3:
|
||||
raise ValueError("Unexpected segmentation output shape")
|
||||
preds = logits_arr.argmax(axis=1)[0]
|
||||
angle_type = get_angle_type(angle_class)
|
||||
if angle_type == "sup":
|
||||
class_map = {
|
||||
0: "background", 1: "effusion", 2: "fat", 3: "fat-pat",
|
||||
4: "femur", 5: "synovium", 6: "tendon",
|
||||
}
|
||||
else:
|
||||
class_map = {
|
||||
0: "background", 1: "fat", 2: "tendon", 3: "muscle",
|
||||
4: "femur", 5: "artery", 6: "baker's cyst",
|
||||
}
|
||||
masks = {}
|
||||
for class_id, class_name in class_map.items():
|
||||
masks[class_name] = (preds == class_id).astype(np.uint8)
|
||||
return preds, masks
|
||||
|
||||
|
||||
async def _get_triton_adapter() -> TritonAdapter:
|
||||
return TritonAdapter(endpoint_url=TRITON_ENDPOINT)
|
||||
|
||||
|
||||
def _encode_image_to_base64(image_pil: Image.Image) -> str:
|
||||
buffered = io.BytesIO()
|
||||
image_pil.save(buffered, format="PNG")
|
||||
return base64.b64encode(buffered.getvalue()).decode()
|
||||
|
||||
|
||||
async def _run_pipeline(image_pil: Image.Image, session_id: str, params: dict, model_versions: dict | None = None) -> dict:
|
||||
enhanced_pil = apply_clahe(image_pil)
|
||||
angle_tensor = prepare_angle_tensor(image_pil)
|
||||
inflammation_tensor = prepare_inflammation_tensor(image_pil)
|
||||
segmentation_tensor = prepare_segmentation_tensor(image_pil)
|
||||
|
||||
triton = await _get_triton_adapter()
|
||||
angle_model = get_model_name("angle", model_versions)
|
||||
angle_result = await triton.infer(
|
||||
model_name=angle_model,
|
||||
inputs={"input": {"data": angle_tensor.tolist(), "shape": list(angle_tensor.shape), "datatype": "FP32"}},
|
||||
)
|
||||
angle_interpreted = _interpret_angle_result(angle_result, params)
|
||||
|
||||
result = {
|
||||
"angle": {
|
||||
"class": angle_interpreted["class"],
|
||||
"confidence": angle_interpreted["confidence"],
|
||||
"calibration": angle_interpreted["calibration"],
|
||||
},
|
||||
"models_used": {"angle": angle_model},
|
||||
}
|
||||
|
||||
if angle_interpreted["class"] in ("post-trans", "sup-up-long"):
|
||||
inflam_model = get_model_name("inflammation", model_versions)
|
||||
inflammation_result = await triton.infer(
|
||||
model_name=inflam_model,
|
||||
inputs={"input": {"data": inflammation_tensor.tolist(), "shape": list(inflammation_tensor.shape), "datatype": "FP32"}},
|
||||
)
|
||||
inflammation_interpreted = _interpret_inflammation_result(inflammation_result, params)
|
||||
result["inflammation"] = {
|
||||
"detected": inflammation_interpreted["detected"],
|
||||
"confidence": inflammation_interpreted["confidence"],
|
||||
"calibration": inflammation_interpreted["calibration"],
|
||||
}
|
||||
result["models_used"]["inflammation"] = inflam_model
|
||||
|
||||
if inflammation_interpreted["detected"]:
|
||||
seg_model_name = get_segmentation_model(angle_interpreted["class"], model_versions)
|
||||
seg_result = await triton.infer(
|
||||
model_name=seg_model_name,
|
||||
inputs={"input": {"data": segmentation_tensor.tolist(), "shape": list(segmentation_tensor.shape), "datatype": "FP32"}},
|
||||
)
|
||||
preds, masks = _process_segmentation_result(seg_result, angle_interpreted["class"])
|
||||
angle_type = get_angle_type(angle_interpreted["class"])
|
||||
measurement = calculate_thickness(masks, image_pil.size)
|
||||
severity = calculate_severity(masks, image_pil.size)
|
||||
segmented_overlay = create_overlay(image_pil, masks, measurement, angle_type)
|
||||
|
||||
result.update({
|
||||
"measurement": measurement,
|
||||
"severity": severity,
|
||||
"segmentation": {
|
||||
"performed": True,
|
||||
"classes_detected": [k for k, v in masks.items() if np.sum(v) > 0],
|
||||
"angle_type": angle_type,
|
||||
},
|
||||
"images": {
|
||||
"enhanced": _encode_image_to_base64(enhanced_pil),
|
||||
"segmented": _encode_image_to_base64(segmented_overlay),
|
||||
},
|
||||
})
|
||||
result["models_used"]["segmentation"] = seg_model_name
|
||||
else:
|
||||
from backend.implementation.pipeline.cv_spec_pipeline import (
|
||||
build_segmentation_skipped,
|
||||
build_severity_zero,
|
||||
)
|
||||
|
||||
result["segmentation"] = build_segmentation_skipped("no_inflammation")
|
||||
result["severity"] = build_severity_zero("no_inflammation")
|
||||
result["images"] = {"enhanced": _encode_image_to_base64(enhanced_pil)}
|
||||
else:
|
||||
from backend.implementation.pipeline.cv_spec_pipeline import (
|
||||
build_segmentation_skipped,
|
||||
build_severity_zero,
|
||||
)
|
||||
|
||||
result["segmentation"] = build_segmentation_skipped("angle_only")
|
||||
result["severity"] = build_severity_zero("angle_only")
|
||||
result["images"] = {"enhanced": _encode_image_to_base64(enhanced_pil)}
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def submit_sync(session_id: str, params: dict, model_versions: dict | None = None) -> JobResult:
|
||||
if "local_image_path" in params:
|
||||
image_pil = Image.open(params["local_image_path"]).convert("RGB")
|
||||
elif "local_image_bytes" in params:
|
||||
image_pil = Image.open(io.BytesIO(params["local_image_bytes"])).convert("RGB")
|
||||
else:
|
||||
from backend.implementation.session import service as session_service
|
||||
frame_metadata = await session_service.get_frame(session_id, params.get("frame_id"))
|
||||
raise NotImplementedError("S3 frame retrieval not yet integrated; use local_image_path for testing")
|
||||
|
||||
pipeline_result = await _run_pipeline(image_pil, session_id, params, model_versions)
|
||||
job_id = str(uuid.uuid4())
|
||||
return JobResult(
|
||||
job_id=job_id,
|
||||
session_id=session_id,
|
||||
status="completed",
|
||||
result=pipeline_result,
|
||||
duration_ms=0,
|
||||
)
|
||||
|
||||
|
||||
async def submit_job(session_id: str, params: dict, model_versions: dict | None = None) -> str:
|
||||
job_id = str(uuid.uuid4())
|
||||
async with _job_lock:
|
||||
_job_registry[job_id] = {
|
||||
"session_id": session_id,
|
||||
"params": params,
|
||||
"model_versions": model_versions,
|
||||
"status": "queued",
|
||||
"result": None,
|
||||
"steps": [],
|
||||
"created_at": datetime.now(),
|
||||
}
|
||||
|
||||
async def _background():
|
||||
try:
|
||||
await push_step_event(job_id, {
|
||||
"step_id": str(uuid.uuid4()),
|
||||
"job_id": job_id,
|
||||
"event_type": "progress",
|
||||
"task_type": "analysis",
|
||||
"status": "running",
|
||||
})
|
||||
async with _job_lock:
|
||||
_job_registry[job_id]["status"] = "running"
|
||||
result = await submit_sync(session_id, params, model_versions)
|
||||
async with _job_lock:
|
||||
_job_registry[job_id].update({
|
||||
"status": "completed",
|
||||
"result": result.model_dump(),
|
||||
})
|
||||
await push_step_event(job_id, {
|
||||
"step_id": str(uuid.uuid4()),
|
||||
"job_id": job_id,
|
||||
"event_type": "completed",
|
||||
"task_type": "analysis",
|
||||
"status": "completed",
|
||||
"data": {"job_result": result.model_dump()},
|
||||
})
|
||||
except Exception as exc:
|
||||
logger.exception(f"Job {job_id} failed")
|
||||
async with _job_lock:
|
||||
_job_registry[job_id].update({
|
||||
"status": "failed",
|
||||
"result": {"error": str(exc)},
|
||||
})
|
||||
await push_step_event(job_id, {
|
||||
"step_id": str(uuid.uuid4()),
|
||||
"job_id": job_id,
|
||||
"event_type": "failed",
|
||||
"task_type": "analysis",
|
||||
"status": "failed",
|
||||
"data": {"error": str(exc)},
|
||||
})
|
||||
|
||||
asyncio.create_task(_background())
|
||||
return job_id
|
||||
|
||||
|
||||
async def job_status(job_id: str) -> JobStatus:
|
||||
async with _job_lock:
|
||||
job = _job_registry.get(job_id)
|
||||
if not job:
|
||||
raise LookupError(f"Job {job_id} not found")
|
||||
return JobStatus(
|
||||
job_id=job_id,
|
||||
session_id=job["session_id"],
|
||||
status=job["status"],
|
||||
result=job.get("result"),
|
||||
steps=job.get("steps", []),
|
||||
created_at=job["created_at"],
|
||||
updated_at=datetime.now(),
|
||||
)
|
||||
|
||||
|
||||
async def job_steps(job_id: str) -> list[PipelineStep]:
|
||||
async with _job_lock:
|
||||
job = _job_registry.get(job_id)
|
||||
if not job:
|
||||
raise LookupError(f"Job {job_id} not found")
|
||||
return job.get("steps", [])
|
||||
|
||||
|
||||
async def list_registered_models() -> ModelCatalog:
|
||||
return ModelCatalog(models=[], total=0)
|
||||
|
||||
|
||||
async def register_model(model_id: str, file: Any) -> ModelRegistrationResult:
|
||||
raise NotImplementedError("Model registration not yet implemented")
|
||||
|
||||
|
||||
async def health() -> HealthStatus:
|
||||
try:
|
||||
triton = await _get_triton_adapter()
|
||||
models = await triton.list_models()
|
||||
ready = any(m.get("state") == "READY" for m in models)
|
||||
status = "ok" if ready else "degraded"
|
||||
return HealthStatus(
|
||||
status=status,
|
||||
version="0.1.0",
|
||||
dependencies={"triton": str(ready)},
|
||||
uptime_seconds=0.0,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(f"Health check failed: {exc}")
|
||||
return HealthStatus(status="error", version="0.1.0", dependencies={"triton": "False"}, uptime_seconds=0.0)
|
||||
|
||||
|
||||
async def push_step_event(job_id: str, event: dict) -> None:
|
||||
from backend.api.analysis_api import _event_queues, _queue_lock
|
||||
async with _queue_lock:
|
||||
if job_id not in _event_queues:
|
||||
_event_queues[job_id] = asyncio.Queue()
|
||||
step_event = StepEvent(
|
||||
step_id=event.get("step_id", ""),
|
||||
job_id=job_id,
|
||||
event_type=event.get("event_type", "progress"),
|
||||
task_type=event.get("task_type", ""),
|
||||
status=event.get("status", "running"),
|
||||
data=event.get("data"),
|
||||
timestamp=datetime.now(),
|
||||
)
|
||||
await _event_queues[job_id].put(step_event)
|
||||
@@ -0,0 +1,21 @@
|
||||
from data.spec.schemas import LoginRequest, Token, UserProfile, UserUpdateRequest
|
||||
|
||||
|
||||
async def login(username: str, password: str) -> Token:
|
||||
raise NotImplementedError("Auth service not yet implemented")
|
||||
|
||||
|
||||
async def logout(token: str) -> None:
|
||||
raise NotImplementedError("Auth service not yet implemented")
|
||||
|
||||
|
||||
async def refresh(refresh_token: str) -> Token:
|
||||
raise NotImplementedError("Auth service not yet implemented")
|
||||
|
||||
|
||||
async def me(token: str) -> UserProfile:
|
||||
raise NotImplementedError("Auth service not yet implemented")
|
||||
|
||||
|
||||
async def update_me(token: str, updates: dict) -> UserProfile:
|
||||
raise NotImplementedError("Auth service not yet implemented")
|
||||
@@ -0,0 +1,75 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
SECRETS_DIR = Path(__file__).resolve().parent.parent.parent.parent.parent.parent / "secrets"
|
||||
|
||||
def _load_secret(name: str, filename: str) -> str:
|
||||
file_path = SECRETS_DIR / filename
|
||||
env_file = os.getenv(f"{name}_FILE")
|
||||
if env_file:
|
||||
resolved = Path(env_file)
|
||||
if resolved.exists():
|
||||
with open(resolved, "r", encoding="utf-8") as f:
|
||||
return f.read().strip()
|
||||
if file_path.exists():
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
return f.read().strip()
|
||||
raise RuntimeError(
|
||||
f"Required secret {name} not found at {file_path} or via {name}_FILE env var"
|
||||
)
|
||||
|
||||
# Endpoints (environment-provided, no hardcoded fallback for production)
|
||||
MODAL_MEDGEMMA_ENDPOINT = os.getenv("MODAL_MEDGEMMA_ENDPOINT")
|
||||
VERTEX_AI_GEMINI_ENDPOINT = os.getenv("VERTEX_AI_GEMINI_ENDPOINT")
|
||||
|
||||
# Secrets (must be present in PILOT_PROJECT/secrets or env)
|
||||
GCP_ACCESS_TOKEN = _load_secret("GCP_ACCESS_TOKEN", "gcp_access_token.txt")
|
||||
MEDGEMMA_API_KEY = _load_secret("MEDGEMMA_API_KEY", "modal_api_key.txt")
|
||||
|
||||
PROJECT_ID = os.getenv("VERTEX_AI_PROJECT", "vkist-project")
|
||||
LOCATION = os.getenv("VERTEX_AI_LOCATION", "asia-southeast1")
|
||||
|
||||
TRITON_ENDPOINT = os.getenv("TRITON_ENDPOINT", "http://localhost:8000")
|
||||
TEMP_DIR = os.getenv("TEMP_DIR", "/tmp/analysis_jobs")
|
||||
|
||||
# LLM Configuration
|
||||
VERTEX_AI_PROJECT = os.getenv("VERTEX_AI_PROJECT", "vkist-project")
|
||||
VERTEX_AI_LOCATION = os.getenv("VERTEX_AI_LOCATION", "asia-southeast1")
|
||||
VERTEX_AI_MODEL = os.getenv("VERTEX_AI_MODEL", "medgemma")
|
||||
|
||||
# Redis Configuration
|
||||
REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
|
||||
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
|
||||
REDIS_DB = int(os.getenv("REDIS_DB", "0"))
|
||||
|
||||
DEFAULT_MODEL_VERSIONS = {
|
||||
"angle": "angle_classify_convnext_tiny",
|
||||
"inflammation": "inflammation_model_efficientnet_b0_ultrasound_2_cls",
|
||||
"segmentation_sup": "segmentation_model_unet_resnet101",
|
||||
"segmentation_post": "segmentation_model_post_deeplabv3_resnet101",
|
||||
}
|
||||
|
||||
CLAHE_CLIP_LIMIT = float(os.getenv("CLAHE_CLIP_LIMIT", "2.0"))
|
||||
CLAHE_TILE_SIZE = tuple(int(x) for x in os.getenv("CLAHE_TILE_SIZE", "8,8").split(","))
|
||||
|
||||
|
||||
def get_model_name(task: str, model_versions: Dict[str, str] | None = None) -> str:
|
||||
if model_versions and task in model_versions:
|
||||
return model_versions[task]
|
||||
return DEFAULT_MODEL_VERSIONS.get(task, task)
|
||||
|
||||
|
||||
def get_angle_type(angle_class: str) -> str:
|
||||
if angle_class in ("sup-trans-flex", "sup-up-long"):
|
||||
return "sup"
|
||||
if angle_class == "post-trans":
|
||||
return "post"
|
||||
return "other"
|
||||
|
||||
|
||||
def get_segmentation_model(angle_class: str, model_versions: Dict[str, str] | None = None) -> str:
|
||||
angle_type = get_angle_type(angle_class)
|
||||
task = "segmentation_sup" if angle_type == "sup" else "segmentation_post"
|
||||
return get_model_name(task, model_versions)
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
from data.spec.schemas import IngestionRecord, RecordDetail
|
||||
from typing import Any
|
||||
|
||||
|
||||
async def list_records(user_id: str) -> list[IngestionRecord]:
|
||||
raise NotImplementedError("Ingestion history service not yet implemented")
|
||||
|
||||
|
||||
async def get_record(record_id: str) -> RecordDetail:
|
||||
raise NotImplementedError("Ingestion history service not yet implemented")
|
||||
@@ -0,0 +1,13 @@
|
||||
from data.spec.schemas import NotificationItem, NotificationPreferences
|
||||
|
||||
|
||||
async def list_notifications(user_id: str, filters: dict | None = None) -> list[NotificationItem]:
|
||||
raise NotImplementedError("Notification service not yet implemented")
|
||||
|
||||
|
||||
async def mark_read(notification_id: str) -> None:
|
||||
raise NotImplementedError("Notification service not yet implemented")
|
||||
|
||||
|
||||
async def set_preferences(user_id: str, prefs: dict) -> None:
|
||||
raise NotImplementedError("Notification service not yet implemented")
|
||||
@@ -0,0 +1,21 @@
|
||||
from data.spec.schemas import Patient, PatientCreate, PatientListResponse
|
||||
|
||||
|
||||
async def list_patients(user_id: str) -> list[Patient]:
|
||||
raise NotImplementedError("Patient service not yet implemented")
|
||||
|
||||
|
||||
async def create_patient(data: dict) -> Patient:
|
||||
raise NotImplementedError("Patient service not yet implemented")
|
||||
|
||||
|
||||
async def get_patient(patient_id: str) -> Patient:
|
||||
raise NotImplementedError("Patient service not yet implemented")
|
||||
|
||||
|
||||
async def list_sessions(patient_id: str) -> list[dict]:
|
||||
raise NotImplementedError("Patient service not yet implemented")
|
||||
|
||||
|
||||
async def ingestion_history(patient_id: str) -> list[dict]:
|
||||
raise NotImplementedError("Patient service not yet implemented")
|
||||
@@ -0,0 +1,13 @@
|
||||
"""CV inference orchestration (Sprint 1–2 spec)."""
|
||||
|
||||
from backend.implementation.pipeline.cv_spec_pipeline import (
|
||||
BRANCH_ANGLE_CLASSES,
|
||||
build_segmentation_skipped,
|
||||
build_severity_zero,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"BRANCH_ANGLE_CLASSES",
|
||||
"build_segmentation_skipped",
|
||||
"build_severity_zero",
|
||||
]
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Shared CV pipeline helpers — Sprint 1–2 architecture spec §7."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# Angles that may run inflammation → conditional segmentation.
|
||||
BRANCH_ANGLE_CLASSES = frozenset({"post-trans", "sup-up-long"})
|
||||
|
||||
|
||||
def build_severity_zero(reason: str) -> dict:
|
||||
descriptions = {
|
||||
"angle_only": "Góc quét không yêu cầu phân đoạn viêm",
|
||||
"no_inflammation": "Không phát hiện viêm — bỏ qua phân đoạn",
|
||||
}
|
||||
return {
|
||||
"level": 0,
|
||||
"severity": "Rất nhẹ",
|
||||
"color": "#28a745",
|
||||
"description": descriptions.get(reason, "Không phân đoạn"),
|
||||
"effusion": {"pixels": 0, "ratio": 0.0, "thickness": 0},
|
||||
"synovium": {"pixels": 0, "ratio": 0.0},
|
||||
"combined_score": 0.0,
|
||||
"reason": reason,
|
||||
}
|
||||
|
||||
|
||||
def build_segmentation_skipped(reason: str) -> dict:
|
||||
notes = {
|
||||
"angle_only": "Chỉ phân loại góc — med-lat / sup-trans-flex",
|
||||
"no_inflammation": "Không phát hiện viêm — bỏ qua phân đoạn",
|
||||
}
|
||||
return {
|
||||
"performed": False,
|
||||
"reason": reason,
|
||||
"note": notes.get(reason, reason),
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
"""Temperature-scaled softmax, entropy guardrails, and risk-first prediction payloads."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
ANGLE_CLASSES = ["med-lat", "post-trans", "sup-trans-flex", "sup-up-long"]
|
||||
INFLAMMATION_CLASSES = ["no_inflammation", "inflammation"]
|
||||
CALIBRATION_TIERS = frozenset({"aggressive", "standard", "conservative"})
|
||||
# Legacy API aliases
|
||||
CALIBRATION_MODES = CALIBRATION_TIERS | frozenset({"screening", "diagnostic"})
|
||||
|
||||
TIER_RECOMMENDED_T = {
|
||||
"aggressive": 0.7,
|
||||
"standard": 1.4,
|
||||
"conservative": 2.2,
|
||||
# legacy → tier
|
||||
"screening": 2.2,
|
||||
"diagnostic": 1.4,
|
||||
}
|
||||
|
||||
TIER_BOUNDARY_AGGRESSIVE_MAX = (0.7 + 1.4) / 2
|
||||
TIER_BOUNDARY_STANDARD_MAX = (1.4 + 2.2) / 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class CalibrationConfig:
|
||||
"""User-adjustable calibration context (maps to UI mode / clinical prior)."""
|
||||
|
||||
temperature: float = 1.4
|
||||
mode: str = "standard"
|
||||
clinical_suspicion: float = 0.0
|
||||
alpha_margin: float = 0.05
|
||||
ood_entropy_threshold: float = 0.88
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.clinical_suspicion = float(np.clip(self.clinical_suspicion, 0.0, 1.0))
|
||||
if self.mode not in CALIBRATION_TIERS:
|
||||
if self.mode in ("screening",):
|
||||
self.mode = "conservative"
|
||||
elif self.mode in ("diagnostic",):
|
||||
self.mode = "standard"
|
||||
else:
|
||||
self.mode = "standard"
|
||||
if self.temperature <= 0:
|
||||
self.temperature = 1.0
|
||||
|
||||
|
||||
def logits_to_array(logits: Any) -> np.ndarray:
|
||||
arr = np.asarray(logits, dtype=np.float32).reshape(-1)
|
||||
if arr.size == 0:
|
||||
raise ValueError("Empty logits")
|
||||
return arr
|
||||
|
||||
|
||||
def resolve_tier_from_temperature(temperature: float) -> str:
|
||||
if temperature <= TIER_BOUNDARY_AGGRESSIVE_MAX:
|
||||
return "aggressive"
|
||||
if temperature <= TIER_BOUNDARY_STANDARD_MAX:
|
||||
return "standard"
|
||||
return "conservative"
|
||||
|
||||
|
||||
def effective_temperature(config: CalibrationConfig) -> float:
|
||||
if config.temperature > 0:
|
||||
return max(0.25, float(config.temperature))
|
||||
return TIER_RECOMMENDED_T.get(config.mode, 1.4)
|
||||
|
||||
|
||||
def temperature_scaled_softmax(logits: np.ndarray, temperature: float) -> np.ndarray:
|
||||
scaled = logits / max(temperature, 1e-6)
|
||||
shifted = scaled - np.max(scaled)
|
||||
exp = np.exp(shifted)
|
||||
return exp / np.sum(exp)
|
||||
|
||||
|
||||
def shannon_entropy(probs: np.ndarray) -> float:
|
||||
safe = np.clip(probs, 1e-12, 1.0)
|
||||
return float(-np.sum(safe * np.log(safe)))
|
||||
|
||||
|
||||
def normalized_entropy(probs: np.ndarray) -> float:
|
||||
if probs.size <= 1:
|
||||
return 0.0
|
||||
return shannon_entropy(probs) / float(np.log(probs.size))
|
||||
|
||||
|
||||
def ambiguous_class_set(probs: np.ndarray, class_names: list[str], alpha_margin: float) -> list[str]:
|
||||
idx_sorted = np.argsort(probs)[::-1]
|
||||
max_prob = float(probs[idx_sorted[0]])
|
||||
return [class_names[i] for i in idx_sorted if float(probs[i]) >= max_prob - alpha_margin]
|
||||
|
||||
|
||||
def estimate_misclassification_rate(max_prob: float) -> float:
|
||||
"""Placeholder empirical mapping until validation-set calibration bins are wired."""
|
||||
if max_prob >= 0.95:
|
||||
return 0.05
|
||||
if max_prob >= 0.90:
|
||||
return 0.08
|
||||
if max_prob >= 0.85:
|
||||
return 0.12
|
||||
if max_prob >= 0.75:
|
||||
return 0.18
|
||||
if max_prob >= 0.65:
|
||||
return 0.25
|
||||
return 0.35
|
||||
|
||||
|
||||
def _risk_framing_vi(
|
||||
predicted_class: str,
|
||||
class_names: list[str],
|
||||
probs: np.ndarray,
|
||||
decision_state: str,
|
||||
error_rate: float,
|
||||
ambiguous_set: list[str],
|
||||
norm_entropy: float,
|
||||
) -> str:
|
||||
if decision_state == "ood_warning":
|
||||
return (
|
||||
"Mô hình chưa được huấn luyện với loại ảnh tương tự, nên kết quả AI có thể không đáng tin. "
|
||||
"Hãy kiểm tra chất lượng ảnh và đối chiếu lâm sàng trước khi dựa vào nhãn tự động."
|
||||
)
|
||||
if decision_state == "ambiguous":
|
||||
alt = ", ".join(c for c in ambiguous_set if c != predicted_class)
|
||||
return (
|
||||
f"Dự đoán chính: {predicted_class}. Tập mơ hồ (α): {', '.join(ambiguous_set)}"
|
||||
+ (f" — các lựa chọn khả dĩ gồm {alt}." if alt else ".")
|
||||
+ " Cần đối chiếu lâm sàng trước khi khóa kết quả."
|
||||
)
|
||||
return (
|
||||
f"Dự đoán: {predicted_class}. "
|
||||
f"Trong các ca có phân bố thống kê tương tự, tỷ lệ phân loại sai ước tính ~{error_rate * 100:.0f}% "
|
||||
f"(entropy chuẩn hóa {norm_entropy:.2f})."
|
||||
)
|
||||
|
||||
|
||||
def decision_state_from(probs: np.ndarray, norm_entropy: float, config: CalibrationConfig) -> str:
|
||||
if norm_entropy >= config.ood_entropy_threshold:
|
||||
return "ood_warning"
|
||||
ambiguous = ambiguous_class_set(probs, [str(i) for i in range(probs.size)], config.alpha_margin)
|
||||
if len(ambiguous) > 1:
|
||||
return "ambiguous"
|
||||
return "confident"
|
||||
|
||||
|
||||
def interpret_classification_logits(
|
||||
logits: Any,
|
||||
class_names: list[str],
|
||||
config: CalibrationConfig | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if len(class_names) == 0:
|
||||
raise ValueError("class_names must not be empty")
|
||||
|
||||
cfg = config or CalibrationConfig()
|
||||
arr = logits_to_array(logits)
|
||||
if arr.size != len(class_names):
|
||||
raise ValueError(f"Expected {len(class_names)} logits, got {arr.size}")
|
||||
|
||||
temperature = effective_temperature(cfg)
|
||||
tier = resolve_tier_from_temperature(temperature)
|
||||
probs = temperature_scaled_softmax(arr, temperature)
|
||||
pred_idx = int(np.argmax(probs))
|
||||
predicted_class = class_names[pred_idx]
|
||||
max_prob = float(probs[pred_idx])
|
||||
entropy = shannon_entropy(probs)
|
||||
norm_entropy = normalized_entropy(probs)
|
||||
ambiguous = ambiguous_class_set(probs, class_names, cfg.alpha_margin)
|
||||
state = decision_state_from(probs, norm_entropy, cfg)
|
||||
error_rate = estimate_misclassification_rate(max_prob)
|
||||
|
||||
risk_vi = _risk_framing_vi(
|
||||
predicted_class,
|
||||
class_names,
|
||||
probs,
|
||||
state,
|
||||
error_rate,
|
||||
ambiguous,
|
||||
norm_entropy,
|
||||
)
|
||||
|
||||
class_probabilities = {
|
||||
name: round(float(probs[i]) * 100, 2) for i, name in enumerate(class_names)
|
||||
}
|
||||
|
||||
return {
|
||||
"class": predicted_class,
|
||||
"confidence": round(max_prob * 100, 2),
|
||||
"calibration": {
|
||||
"raw_logits": [round(float(x), 6) for x in arr.tolist()],
|
||||
"temperature": round(temperature, 4),
|
||||
"base_temperature": cfg.temperature,
|
||||
"mode": tier,
|
||||
"clinical_suspicion": round(cfg.clinical_suspicion, 3),
|
||||
"alpha_margin": cfg.alpha_margin,
|
||||
"class_probabilities": class_probabilities,
|
||||
"entropy": round(entropy, 4),
|
||||
"normalized_entropy": round(norm_entropy, 4),
|
||||
"ambiguous_set": ambiguous,
|
||||
"decision_state": state,
|
||||
"predicted_error_rate": round(error_rate, 4),
|
||||
"risk_framing_vi": risk_vi,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def interpret_angle_logits(logits: Any, config: CalibrationConfig | None = None) -> dict[str, Any]:
|
||||
return interpret_classification_logits(logits, ANGLE_CLASSES, config)
|
||||
|
||||
|
||||
def interpret_inflammation_logits(logits: Any, config: CalibrationConfig | None = None) -> dict[str, Any]:
|
||||
payload = interpret_classification_logits(logits, INFLAMMATION_CLASSES, config)
|
||||
detected = payload["class"] == "inflammation"
|
||||
payload["detected"] = detected
|
||||
return payload
|
||||
|
||||
|
||||
def calibration_config_from_params(params: dict[str, Any] | None) -> CalibrationConfig:
|
||||
if not params:
|
||||
return CalibrationConfig()
|
||||
calibration = params.get("calibration") or {}
|
||||
return CalibrationConfig(
|
||||
temperature=float(calibration.get("temperature", 1.0)),
|
||||
mode=str(calibration.get("mode", "standard")),
|
||||
clinical_suspicion=float(calibration.get("clinical_suspicion", 0.0)),
|
||||
alpha_margin=float(calibration.get("alpha_margin", 0.05)),
|
||||
ood_entropy_threshold=float(calibration.get("ood_entropy_threshold", 0.88)),
|
||||
)
|
||||
@@ -0,0 +1,123 @@
|
||||
__all__ = ["calculate_thickness", "get_mask_bounding_box", "find_max_continuous_segment"]
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
SEGMENT_CLASSES_SUPRAPAT = {
|
||||
0: "background", 1: "effusion", 2: "fat", 3: "fat-pat",
|
||||
4: "femur", 5: "synovium", 6: "tendon"
|
||||
}
|
||||
SEGMENT_CLASSES_POST = {
|
||||
0: "background", 1: "fat", 2: "tendon", 3: "muscle",
|
||||
4: "femur", 5: "artery", 6: "baker's cyst"
|
||||
}
|
||||
PIXEL_TO_MM = 45.0 / 655.0
|
||||
|
||||
|
||||
def get_mask_bounding_box(mask, dist_percent: float = 0.01):
|
||||
if mask is None or np.sum(mask) == 0:
|
||||
return None
|
||||
mask_uint8 = mask.astype(np.uint8)
|
||||
if np.max(mask_uint8) == 1:
|
||||
mask_uint8 *= 255
|
||||
img_width = mask_uint8.shape[1]
|
||||
dist_threshold = img_width * dist_percent
|
||||
kernel = np.ones((5, 5), np.uint8)
|
||||
clean_mask = cv2.morphologyEx(mask_uint8, cv2.MORPH_OPEN, kernel)
|
||||
contours, _ = cv2.findContours(clean_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
if not contours:
|
||||
return None
|
||||
contour_info = sorted(
|
||||
[{"cnt": cnt, "area": cv2.contourArea(cnt)} for cnt in contours],
|
||||
key=lambda x: x["area"], reverse=True,
|
||||
)
|
||||
main_block = contour_info[0]
|
||||
max_area = main_block["area"]
|
||||
if max_area < 50:
|
||||
return None
|
||||
main_mask = np.zeros_like(mask_uint8)
|
||||
cv2.drawContours(main_mask, [main_block["cnt"]], -1, 255, -1)
|
||||
dist_map = cv2.distanceTransform(255 - main_mask, cv2.DIST_L2, 3)
|
||||
significant_contours = [main_block["cnt"]]
|
||||
area_threshold = max_area / 4.0
|
||||
for i in range(1, len(contour_info)):
|
||||
other = contour_info[i]
|
||||
other_mask = np.zeros_like(mask_uint8)
|
||||
cv2.drawContours(other_mask, [other["cnt"]], -1, 255, -1)
|
||||
min_dist = np.min(dist_map[other_mask > 0])
|
||||
if other["area"] >= area_threshold or min_dist <= dist_threshold:
|
||||
significant_contours.append(other["cnt"])
|
||||
all_points = np.concatenate(significant_contours)
|
||||
x, y, w, h = cv2.boundingRect(all_points)
|
||||
return x, y, w, h
|
||||
|
||||
|
||||
def find_max_continuous_segment(col_array):
|
||||
padded = np.concatenate(([0], col_array, [0]))
|
||||
diffs = np.diff(padded)
|
||||
starts = np.where(diffs == 1)[0]
|
||||
ends = np.where(diffs == -1)[0]
|
||||
if len(starts) == 0:
|
||||
return 0, -1, -1
|
||||
lengths = ends - starts
|
||||
max_idx = int(np.argmax(lengths))
|
||||
max_len = int(lengths[max_idx])
|
||||
return max_len, int(starts[max_idx]), int(ends[max_idx])
|
||||
|
||||
|
||||
def calculate_thickness(masks: dict, image_size, measure_ids=None):
|
||||
if measure_ids is None:
|
||||
measure_ids = [1, 5]
|
||||
width, height = image_size
|
||||
mask_all_labels = np.zeros((height, width), dtype=np.uint8)
|
||||
mask_measure = np.zeros((height, width), dtype=np.uint8)
|
||||
has_any_label = False
|
||||
if "fat-pat" in masks:
|
||||
class_map = SEGMENT_CLASSES_SUPRAPAT
|
||||
else:
|
||||
class_map = SEGMENT_CLASSES_POST
|
||||
for class_id, class_name in class_map.items():
|
||||
if class_name not in masks or class_name == "background":
|
||||
continue
|
||||
mask = masks[class_name]
|
||||
if np.sum(mask) > 0:
|
||||
has_any_label = True
|
||||
mask_all_labels = np.logical_or(mask_all_labels, mask).astype(np.uint8)
|
||||
if class_id in measure_ids:
|
||||
mask_measure = np.logical_or(mask_measure, mask).astype(np.uint8)
|
||||
if not has_any_label or np.sum(mask_measure) == 0:
|
||||
return None
|
||||
bbox_all = get_mask_bounding_box(mask_all_labels)
|
||||
if bbox_all is None:
|
||||
return None
|
||||
x_all, y_all, w_all, h_all = bbox_all
|
||||
roi_start = x_all + (w_all // 3)
|
||||
roi_end = x_all + (2 * w_all // 3)
|
||||
roi_strip = mask_measure[:, roi_start:roi_end]
|
||||
global_max_len_px = 0
|
||||
best_x_rel = 0
|
||||
best_y_start = 0
|
||||
best_y_end = 0
|
||||
for x in range(roi_strip.shape[1]):
|
||||
col = roi_strip[:, x]
|
||||
if not np.any(col):
|
||||
continue
|
||||
length, y_s, y_e = find_max_continuous_segment(col)
|
||||
if length > global_max_len_px:
|
||||
global_max_len_px = length
|
||||
best_x_rel = x
|
||||
best_y_start = y_s
|
||||
best_y_end = y_e
|
||||
if global_max_len_px == 0:
|
||||
return None
|
||||
thickness_mm = global_max_len_px * PIXEL_TO_MM
|
||||
real_x = roi_start + best_x_rel
|
||||
return {
|
||||
"thickness_px": int(global_max_len_px),
|
||||
"thickness_mm": float(round(thickness_mm, 2)),
|
||||
"x": int(real_x),
|
||||
"y_start": int(best_y_start),
|
||||
"y_end": int(best_y_end),
|
||||
"roi_start": int(roi_start),
|
||||
"roi_end": int(roi_end),
|
||||
"bbox": {"x": int(x_all), "y": int(y_all), "w": int(w_all), "h": int(h_all)},
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
__all__ = ["create_overlay"]
|
||||
from PIL import Image, ImageDraw
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
COLOR_MAP_SUP = {
|
||||
"background": [0, 0, 0],
|
||||
"effusion": [255, 0, 0],
|
||||
"fat": [255, 255, 0],
|
||||
"fat-pat": [0, 255, 255],
|
||||
"femur": [0, 255, 0],
|
||||
"synovium": [255, 0, 255],
|
||||
"tendon": [0, 0, 255],
|
||||
}
|
||||
|
||||
COLOR_MAP_POST = {
|
||||
"background": [0, 0, 0],
|
||||
"baker's cyst": [255, 0, 0],
|
||||
"fat": [255, 255, 0],
|
||||
"muscle": [0, 255, 255],
|
||||
"femur": [0, 255, 0],
|
||||
"artery": [255, 0, 255],
|
||||
"synovium": [255, 0, 255],
|
||||
"tendon": [0, 0, 255],
|
||||
}
|
||||
|
||||
|
||||
def create_overlay(image_pil: Image.Image, masks: dict, measurement, angle_type: str = "sup") -> Image.Image:
|
||||
if masks is None:
|
||||
return image_pil
|
||||
color_map = COLOR_MAP_SUP if angle_type == "sup" else COLOR_MAP_POST
|
||||
img_array = np.array(image_pil)
|
||||
overlay = img_array.copy()
|
||||
for class_name, mask in masks.items():
|
||||
if class_name in color_map and np.sum(mask) > 0:
|
||||
color = color_map[class_name]
|
||||
for i in range(3):
|
||||
overlay[:, :, i] = np.where(
|
||||
mask > 0,
|
||||
(overlay[:, :, i] * 0.6 + color[i] * 0.4).astype(np.uint8),
|
||||
overlay[:, :, i],
|
||||
)
|
||||
overlay_pil = Image.fromarray(overlay)
|
||||
draw = ImageDraw.Draw(overlay_pil)
|
||||
for class_name in ["effusion", "synovium"]:
|
||||
mask = masks.get(class_name)
|
||||
if mask is not None and np.sum(mask) > 0:
|
||||
mask_uint8 = (mask * 255).astype(np.uint8)
|
||||
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
for contour in contours:
|
||||
points = contour.reshape(-1, 2).tolist()
|
||||
if len(points) > 2:
|
||||
points = [(int(p[0]), int(p[1])) for p in points]
|
||||
draw.line(points + [points[0]], fill=(255, 255, 255), width=3)
|
||||
if measurement and angle_type == "sup":
|
||||
x = measurement["x"]
|
||||
y_start = measurement["y_start"]
|
||||
y_end = measurement["y_end"]
|
||||
thickness_mm = measurement["thickness_mm"]
|
||||
roi_start = measurement["roi_start"]
|
||||
roi_end = measurement["roi_end"]
|
||||
bbox = measurement["bbox"]
|
||||
draw.rectangle(
|
||||
[bbox["x"], bbox["y"], bbox["x"] + bbox["w"], bbox["y"] + bbox["h"]],
|
||||
outline=(0, 255, 0), width=3,
|
||||
)
|
||||
h = image_pil.size[1]
|
||||
draw.line([(roi_start, 0), (roi_start, h)], fill=(0, 255, 255), width=2)
|
||||
draw.line([(roi_end, 0), (roi_end, h)], fill=(0, 255, 255), width=2)
|
||||
draw.line([(x, y_start), (x, y_end)], fill=(255, 0, 0), width=4)
|
||||
radius = 4
|
||||
draw.ellipse([x - radius, y_start - radius, x + radius, y_start + radius],
|
||||
fill=(0, 255, 0), outline=(255, 255, 255), width=2)
|
||||
draw.ellipse([x - radius, y_end - radius, x + radius, y_end + radius],
|
||||
fill=(0, 255, 0), outline=(255, 255, 255), width=2)
|
||||
text = f"{thickness_mm:.2f} mm"
|
||||
try:
|
||||
from PIL import ImageFont
|
||||
font = ImageFont.load_default()
|
||||
bbox_text = draw.textbbox((0, 0), text, font=font)
|
||||
text_w = bbox_text[2] - bbox_text[0]
|
||||
text_h = bbox_text[3] - bbox_text[1]
|
||||
except Exception:
|
||||
text_w, text_h = 100, 20
|
||||
text_x = x + 8
|
||||
text_y = y_start - text_h - 8
|
||||
draw.rectangle([text_x - 2, text_y - 2, text_x + text_w + 2, text_y + text_h + 2], fill=(0, 0, 0))
|
||||
draw.text((text_x, text_y), text, fill=(255, 0, 0))
|
||||
return overlay_pil
|
||||
@@ -0,0 +1,64 @@
|
||||
__all__ = ["calculate_severity"]
|
||||
import numpy as np
|
||||
|
||||
SEVERITY_LEVELS = [
|
||||
(15, 3, "Nặng", "#dc3545", "Dịch khớp dày, màng hoạt dịch tăng sinh rõ"),
|
||||
(8, 2, "Trung bình", "#fd7e14", "Dịch khớp trung bình, màng hoạt dịch tăng sinh vừa"),
|
||||
(3, 1, "Nhẹ", "#ffc107", "Dịch khớp mỏng, màng hoạt dịch tăng sinh nhẹ"),
|
||||
(0, 0, "Rất nhẹ", "#28a745", "Lượng dịch và màng hoạt dịch trong giới hạn bình thường"),
|
||||
]
|
||||
|
||||
|
||||
def calculate_severity(masks: dict, image_size) -> dict | None:
|
||||
if not masks:
|
||||
return None
|
||||
width, height = image_size
|
||||
total_pixels = width * height
|
||||
effusion_mask = masks.get("effusion", np.zeros((height, width), dtype=np.uint8))
|
||||
effusion_pixels = int(np.sum(effusion_mask))
|
||||
effusion_ratio = (effusion_pixels / total_pixels) * 100
|
||||
effusion_thickness = 0
|
||||
if effusion_pixels > 0:
|
||||
rows_with_effusion = np.any(effusion_mask > 0, axis=1)
|
||||
if np.any(rows_with_effusion):
|
||||
effusion_thickness = int(np.sum(rows_with_effusion))
|
||||
synovium_mask = masks.get("synovium", np.zeros((height, width), dtype=np.uint8))
|
||||
synovium_pixels = int(np.sum(synovium_mask))
|
||||
synovium_ratio = (synovium_pixels / total_pixels) * 100
|
||||
effusion_score = min(effusion_thickness / height * 100, 100)
|
||||
synovium_score = synovium_ratio
|
||||
combined_score = effusion_score * 0.6 + synovium_score * 0.4
|
||||
for threshold, level, severity, color, description in SEVERITY_LEVELS:
|
||||
if combined_score > threshold:
|
||||
return {
|
||||
"level": int(level),
|
||||
"severity": severity,
|
||||
"color": color,
|
||||
"description": description,
|
||||
"effusion": {
|
||||
"pixels": effusion_pixels,
|
||||
"ratio": float(round(effusion_ratio, 2)),
|
||||
"thickness": effusion_thickness,
|
||||
},
|
||||
"synovium": {
|
||||
"pixels": synovium_pixels,
|
||||
"ratio": float(round(synovium_ratio, 2)),
|
||||
},
|
||||
"combined_score": float(round(combined_score, 2)),
|
||||
}
|
||||
return {
|
||||
"level": 0,
|
||||
"severity": "Rất nhẹ",
|
||||
"color": "#28a745",
|
||||
"description": "Lượng dịch và màng hoạt dịch trong giới hạn bình thường",
|
||||
"effusion": {
|
||||
"pixels": effusion_pixels,
|
||||
"ratio": float(round(effusion_ratio, 2)),
|
||||
"thickness": effusion_thickness,
|
||||
},
|
||||
"synovium": {
|
||||
"pixels": synovium_pixels,
|
||||
"ratio": float(round(synovium_ratio, 2)),
|
||||
},
|
||||
"combined_score": float(round(combined_score, 2)),
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
__all__ = ["apply_clahe"]
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def apply_clahe(image_pil: Image.Image, clip_limit: float = 2.0, tile_grid_size: tuple[int, int] = (8, 8)) -> Image.Image:
|
||||
img_array = np.array(image_pil)
|
||||
gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
|
||||
clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
|
||||
enhanced_gray = clahe.apply(gray)
|
||||
enhanced_rgb = cv2.cvtColor(enhanced_gray, cv2.COLOR_GRAY2RGB)
|
||||
return Image.fromarray(enhanced_rgb)
|
||||
@@ -0,0 +1,35 @@
|
||||
__all__ = ["prepare_angle_tensor", "prepare_inflammation_tensor", "prepare_segmentation_tensor"]
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from .transforms import Resize, Normalize
|
||||
|
||||
ANGLE_TRANSFORM = Resize((224, 224))
|
||||
ANGLE_NORMALIZE = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
||||
|
||||
INFLAMMATION_TRANSFORM = Resize((224, 224))
|
||||
INFLAMMATION_NORMALIZE = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
||||
|
||||
SEGMENTATION_TRANSFORM = Resize((512, 512))
|
||||
|
||||
|
||||
def _to_nchw(arr_hwc: np.ndarray) -> np.ndarray:
|
||||
arr = arr_hwc.transpose(2, 0, 1)
|
||||
return np.expand_dims(arr, axis=0)
|
||||
|
||||
|
||||
def prepare_angle_tensor(image_pil: Image.Image) -> np.ndarray:
|
||||
img = ANGLE_TRANSFORM(image_pil)
|
||||
arr = ANGLE_NORMALIZE(img)
|
||||
return _to_nchw(arr)
|
||||
|
||||
|
||||
def prepare_inflammation_tensor(image_pil: Image.Image) -> np.ndarray:
|
||||
img = INFLAMMATION_TRANSFORM(image_pil)
|
||||
arr = INFLAMMATION_NORMALIZE(img)
|
||||
return _to_nchw(arr)
|
||||
|
||||
|
||||
def prepare_segmentation_tensor(image_pil: Image.Image) -> np.ndarray:
|
||||
img = SEGMENTATION_TRANSFORM(image_pil)
|
||||
arr = np.asarray(img).astype(np.float32) / 255.0
|
||||
return _to_nchw(arr)
|
||||
@@ -0,0 +1,22 @@
|
||||
__all__ = ["Resize", "Normalize"]
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Resize:
|
||||
def __init__(self, size: tuple[int, int]):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, image: Image.Image) -> Image.Image:
|
||||
return image.resize(self.size, Image.Resampling.BILINEAR)
|
||||
|
||||
|
||||
class Normalize:
|
||||
def __init__(self, mean: list[float], std: list[float]):
|
||||
self.mean = np.array(mean, dtype=np.float32)
|
||||
self.std = np.array(std, dtype=np.float32)
|
||||
|
||||
def __call__(self, image_pil: Image.Image) -> np.ndarray:
|
||||
arr = np.asarray(image_pil).astype(np.float32) / 255.0
|
||||
arr = (arr - self.mean) / self.std
|
||||
return arr
|
||||
@@ -0,0 +1,148 @@
|
||||
from typing import Any, AsyncGenerator
|
||||
import logging
|
||||
from fastapi import HTTPException, status
|
||||
from data.spec.schemas import (
|
||||
HeatmapResult, RationaleResult, ChatResponse, DriftCheckResult,
|
||||
EvidenceList, ActivationMeta, AnnotationArtifact, EscalationTicket,
|
||||
GuardrailResult, CorrectionRecord,
|
||||
)
|
||||
from backend.implementation.adapters.llm_adapter import get_llm_adapter
|
||||
from backend.implementation.adapters.bert_adapter import get_bert_adapter
|
||||
from backend.implementation.adapters.redis_adapter import get_redis_client
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
llm_adapter = get_llm_adapter()
|
||||
bert_adapter = get_bert_adapter()
|
||||
redis_client = get_redis_client()
|
||||
|
||||
async def _verify_pre_egress(session_id: str, redaction_hash: str | None = None):
|
||||
"""
|
||||
Enforce NFR-16a Pre-Egress Checklist.
|
||||
"""
|
||||
# 1. Consent Verification
|
||||
consent_key = f"consent:{session_id}"
|
||||
if not redis_client.exists(consent_key):
|
||||
logger.error(f"Consent missing for session {session_id}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="User consent for cloud LLM egress is required."
|
||||
)
|
||||
|
||||
# 2. Redaction Verification (if hash provided)
|
||||
if redaction_hash:
|
||||
# In real impl: Run Presidio on the prompt and compare hashes
|
||||
# For now, we assume a simple check or stub
|
||||
if redaction_hash == "FAIL_HASH":
|
||||
logger.error(f"Redaction hash mismatch for session {session_id}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Redaction verification failed. PHI may be present."
|
||||
)
|
||||
|
||||
# Note: Audit Log commit is handled via the LLM adapter's AuditCallbackHandler
|
||||
# to ensure it happens exactly before the call.
|
||||
|
||||
async def gradcam(session_id: str) -> HeatmapResult:
|
||||
raise NotImplementedError("Safety service not yet implemented")
|
||||
|
||||
|
||||
async def rationale(session_id: str, redaction_hash: str | None = None) -> RationaleResult:
|
||||
# Pre-egress check
|
||||
await _verify_pre_egress(session_id, redaction_hash)
|
||||
|
||||
# 1. Fetch session context (simplified for stub)
|
||||
context = {"grade": "moderate", "joint_site": "wrist"}
|
||||
|
||||
# 2. Construct prompt
|
||||
prompt = f"Based on MOH guidelines, explain the synovitis grade {context['grade']} for {context['joint_site']}..."
|
||||
|
||||
# 3. Call LLM adapter
|
||||
text = await llm_adapter.generate(prompt, session_id)
|
||||
|
||||
redis_client.set(f"consult_mode:{session_id}", "tier_3")
|
||||
|
||||
return RationaleResult(text=text)
|
||||
|
||||
|
||||
async def circuit_break(session_id: str, flag: bool) -> None:
|
||||
if flag:
|
||||
logger.warning(f"Circuit breaker triggered for session {session_id}")
|
||||
return None
|
||||
|
||||
|
||||
async def socratic_chat(session_id: str, prompt: str, redaction_hash: str | None = None) -> ChatResponse:
|
||||
# Pre-egress check
|
||||
await _verify_pre_egress(session_id, redaction_hash)
|
||||
|
||||
# 1. Retrieve conversation history (stub)
|
||||
history = []
|
||||
|
||||
# 2. Construct prompt
|
||||
full_prompt = f"History: {history}\nUser: {prompt}\nAssistant: "
|
||||
|
||||
# 3. Call LLM adapter
|
||||
response_text = await llm_adapter.generate(full_prompt, session_id)
|
||||
|
||||
# 4. BERT Referee check (stub)
|
||||
is_grounded = bert_adapter.referee_check(response_text, [])
|
||||
if not is_grounded:
|
||||
response_text = "I'm sorry, I couldn't verify this answer against the guidelines."
|
||||
|
||||
# Post-egress: update consult mode
|
||||
redis_client.set(f"consult_mode:{session_id}", "tier_3")
|
||||
|
||||
return ChatResponse(response=response_text)
|
||||
|
||||
|
||||
async def drift_check(session_id: str) -> DriftCheckResult:
|
||||
res = bert_adapter.drift_check("mock clinical text")
|
||||
return DriftCheckResult(score=res.score, is_drifted=res.is_drifted)
|
||||
|
||||
|
||||
async def rag_evidence(session_id: str) -> EvidenceList:
|
||||
raise NotImplementedError("Safety service not yet implemented")
|
||||
|
||||
|
||||
async def activations(session_id: str, params: dict) -> ActivationMeta:
|
||||
raise NotImplementedError("Safety service not yet implemented")
|
||||
|
||||
|
||||
async def upload_artifact(session_id: str, file: Any) -> AnnotationArtifact:
|
||||
raise NotImplementedError("Safety service not yet implemented")
|
||||
|
||||
|
||||
async def ground_truth(session_id: str, label: dict) -> None:
|
||||
raise NotImplementedError("Safety service not yet implemented")
|
||||
|
||||
|
||||
async def escalate(session_id: str, reason: str) -> EscalationTicket:
|
||||
raise NotImplementedError("Safety service not yet implemented")
|
||||
|
||||
|
||||
async def morphology(session_id: str, annotation: dict) -> None:
|
||||
raise NotImplementedError("Safety service not yet implemented")
|
||||
|
||||
|
||||
async def guardrail_check(session_id: str, prompt: str, score: float) -> GuardrailResult:
|
||||
res = bert_adapter.guardrail_check(prompt)
|
||||
return GuardrailResult(verdict=res.verdict, reason=res.reason)
|
||||
|
||||
|
||||
async def submit_correction(session_id: str, correction: dict) -> CorrectionRecord:
|
||||
raise NotImplementedError("Safety correction service not yet implemented")
|
||||
|
||||
|
||||
async def chat_stream(session_id: str, prompt: str, redaction_hash: str | None = None) -> AsyncGenerator[str, None]:
|
||||
# Pre-egress check
|
||||
await _verify_pre_egress(session_id, redaction_hash)
|
||||
|
||||
async for chunk in llm_adapter.stream_generate(prompt, session_id):
|
||||
res = bert_adapter.guardrail_check(chunk)
|
||||
if res.verdict == "MITIGATE":
|
||||
yield "[Content Filtered]"
|
||||
return
|
||||
yield chunk
|
||||
|
||||
# Post-egress
|
||||
redis_client.set(f"consult_mode:{session_id}", "tier_3")
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
from datetime import datetime
|
||||
from data.spec.schemas import (
|
||||
Session, SessionCreate, SessionDetail, SessionPatchReview,
|
||||
FrameMetadata, PersistResult, ExportResult, ScrubResult,
|
||||
)
|
||||
from typing import Any
|
||||
|
||||
|
||||
async def create_session(user_id: str, patient_id: str, case_id: str | None = None) -> Session:
|
||||
raise NotImplementedError("Session service not yet implemented")
|
||||
|
||||
|
||||
async def get_session(session_id: str) -> SessionDetail:
|
||||
raise NotImplementedError("Session service not yet implemented")
|
||||
|
||||
|
||||
async def add_frame(session_id: str, file: Any, frame_number: int | None = None) -> FrameMetadata:
|
||||
raise NotImplementedError("Session service not yet implemented")
|
||||
|
||||
|
||||
async def patch_review(session_id: str, review: dict) -> Session:
|
||||
raise NotImplementedError("Session service not yet implemented")
|
||||
|
||||
|
||||
async def persist(session_id: str, review: dict) -> PersistResult:
|
||||
raise NotImplementedError("Session service not yet implemented")
|
||||
|
||||
|
||||
async def export_pdf(session_id: str, params: dict) -> ExportResult:
|
||||
raise NotImplementedError("Session service not yet implemented")
|
||||
|
||||
|
||||
async def scrub_validate(session_id: str, metadata: dict) -> ScrubResult:
|
||||
raise NotImplementedError("Session service not yet implemented")
|
||||
@@ -0,0 +1,9 @@
|
||||
from data.spec.schemas import UserSettings, SettingsUpdate
|
||||
|
||||
|
||||
async def get_settings(user_id: str) -> UserSettings:
|
||||
raise NotImplementedError("Settings service not yet implemented")
|
||||
|
||||
|
||||
async def update_settings(user_id: str, updates: dict) -> UserSettings:
|
||||
raise NotImplementedError("Settings service not yet implemented")
|
||||
@@ -0,0 +1,6 @@
|
||||
from data.spec.schemas import AnomalyRecord
|
||||
from typing import Any
|
||||
|
||||
|
||||
async def report_anomaly(session_id: str, data: dict) -> AnomalyRecord:
|
||||
raise NotImplementedError("Telemetry service not yet implemented")
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Triton batching helpers — aligned with config.pbtxt max_batch_size: 8."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Iterator, Sequence
|
||||
from typing import TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
TRITON_MAX_BATCH_SIZE = int(os.getenv("TRITON_MAX_BATCH_SIZE", "8"))
|
||||
|
||||
|
||||
def chunk_sequence(items: Sequence[T], batch_size: int | None = None) -> Iterator[list[T]]:
|
||||
"""Split a sequence into chunks of at most ``batch_size`` (default: TRITON_MAX_BATCH_SIZE)."""
|
||||
size = batch_size if batch_size is not None else TRITON_MAX_BATCH_SIZE
|
||||
if size < 1:
|
||||
raise ValueError(f"batch_size must be >= 1, got {size}")
|
||||
for start in range(0, len(items), size):
|
||||
yield list(items[start : start + size])
|
||||
|
||||
|
||||
def batch_count(item_count: int, batch_size: int | None = None) -> int:
|
||||
"""Number of Triton infer calls needed (e.g. 10 images -> 2 batches when size=8)."""
|
||||
if item_count <= 0:
|
||||
return 0
|
||||
size = batch_size if batch_size is not None else TRITON_MAX_BATCH_SIZE
|
||||
return (item_count + size - 1) // size
|
||||
Reference in New Issue
Block a user