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')
|
||||
def detect_objects(photo_id: str):
|
||||
"""Detect objects in a photo — implemented in PR6."""
|
||||
return {'status': 'not_implemented'}
|
||||
"""Detect objects in a photo, create Tag(kind=object) rows, and
|
||||
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')
|
||||
|
||||
Reference in New Issue
Block a user