feat: add YOLO object detection writing to unified Tag model
Implement detect_objects Celery task: - Runs YOLOv8n on 640px thumbnail via ONNX Runtime - Creates Tag(kind=object) rows for each COCO class detected - Writes photo_tags associations with confidence, bbox, and source - Wipes previous detections per source model on re-run No new tables/migrations — uses the unified Tag model from PR3. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -134,8 +134,66 @@ async def _ocr_photo_async(photo_id: str):
|
|||||||
|
|
||||||
@shared_task(name='detect_objects', queue='vision')
|
@shared_task(name='detect_objects', queue='vision')
|
||||||
def detect_objects(photo_id: str):
|
def detect_objects(photo_id: str):
|
||||||
"""Detect objects in a photo — implemented in PR6."""
|
"""Detect objects in a photo, create Tag(kind=object) rows, and
|
||||||
return {'status': 'not_implemented'}
|
link via photo_tags with confidence/bbox/source."""
|
||||||
|
if not settings.vision.enabled or not settings.vision.detector.enabled:
|
||||||
|
return {'status': 'skipped', 'reason': 'detection disabled'}
|
||||||
|
return asyncio.run(_detect_objects_async(photo_id))
|
||||||
|
|
||||||
|
|
||||||
|
async def _detect_objects_async(photo_id: str):
|
||||||
|
image = _load_thumb(photo_id, "medium") # 640px
|
||||||
|
if image is None:
|
||||||
|
return {'status': 'error', 'message': 'thumbnail not found'}
|
||||||
|
|
||||||
|
from app.services.vision.registry import registry
|
||||||
|
detector = registry.get_detector()
|
||||||
|
detections = detector.detect(image)
|
||||||
|
|
||||||
|
if not detections:
|
||||||
|
logger.info("No objects detected for photo %s", photo_id)
|
||||||
|
return {'status': 'success', 'photo_id': photo_id, 'objects': 0}
|
||||||
|
|
||||||
|
from app.models.tags import Tag, photo_tags
|
||||||
|
|
||||||
|
source_name = "vision:yolov8n"
|
||||||
|
|
||||||
|
async with AsyncSessionLocal() as session:
|
||||||
|
# Wipe previous detection results for this photo from this model
|
||||||
|
await session.execute(
|
||||||
|
delete(photo_tags).where(
|
||||||
|
photo_tags.c.photo_id == photo_id,
|
||||||
|
photo_tags.c.source == source_name,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
for det in detections:
|
||||||
|
# Find or create the object tag
|
||||||
|
result = await session.execute(
|
||||||
|
select(Tag).where(Tag.name == det.label, Tag.kind == 'object')
|
||||||
|
)
|
||||||
|
tag = result.scalar_one_or_none()
|
||||||
|
if not tag:
|
||||||
|
tag = Tag(name=det.label, kind='object', source=source_name)
|
||||||
|
session.add(tag)
|
||||||
|
await session.flush() # get tag.id
|
||||||
|
|
||||||
|
# Insert photo_tags association with ML metadata
|
||||||
|
await session.execute(
|
||||||
|
photo_tags.insert().values(
|
||||||
|
photo_id=photo_id,
|
||||||
|
tag_id=tag.id,
|
||||||
|
confidence=det.confidence,
|
||||||
|
bbox=det.bbox,
|
||||||
|
source=source_name,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
await session.commit()
|
||||||
|
|
||||||
|
labels = [d.label for d in detections]
|
||||||
|
logger.info("Detected %d objects in photo %s: %s", len(detections), photo_id, labels)
|
||||||
|
return {'status': 'success', 'photo_id': photo_id, 'objects': len(detections)}
|
||||||
|
|
||||||
|
|
||||||
@shared_task(name='extract_faces', queue='vision')
|
@shared_task(name='extract_faces', queue='vision')
|
||||||
|
|||||||
Reference in New Issue
Block a user