model_train_dm/backend/apps/common/utils/annotation_store.py
2026-07-27 17:51:49 +08:00

62 lines
1.8 KiB
Python

import json
from django.db.models import TextField, Value
from django.db.models.functions import Coalesce
def _safe_json_loads(raw):
if raw in (None, ""):
return None
try:
return json.loads(raw)
except Exception:
return None
def is_polygon_annotation_content(content) -> bool:
parsed = _safe_json_loads(content)
if isinstance(parsed, dict):
if parsed.get("type") == "polygon":
return True
if isinstance(parsed.get("shapes"), list):
return True
if isinstance(parsed, list):
for item in parsed:
if not isinstance(item, dict):
continue
region = item.get("region")
if isinstance(region, list) and region and isinstance(region[0], (list, tuple)):
return True
return False
def split_dataset_sample_annotation_content(content):
raw = content or ""
if not raw:
return {"rectangle": None, "polygon": None}
if is_polygon_annotation_content(raw):
return {"rectangle": None, "polygon": raw}
return {"rectangle": raw, "polygon": None}
def get_dataset_sample_annotation_content(sample) -> str:
if sample is None:
return ""
polygon = getattr(sample, "polygon", None)
rectangle = getattr(sample, "rectangle", None)
return polygon or rectangle or ""
def get_dataset_sample_annotation_from_row(sample_row: dict) -> str:
if not isinstance(sample_row, dict):
return ""
return sample_row.get("polygon") or sample_row.get("rectangle") or sample_row.get("annotation_content") or ""
def annotate_queryset_with_annotation_content(queryset, alias: str = "annotation_content"):
return queryset.annotate(
**{
alias: Coalesce("polygon", "rectangle", Value("", output_field=TextField()), output_field=TextField()),
}
)