wrongmove/services/floorplan_detector.py

"""Floorplan detector service - OCR-based square meter detection."""
import asyncio
import logging
from config.scraper_config import MAX_OCR_WORKERS
from models import Listing
from rec import floorplan
from repositories.listing_repository import ListingRepository

logger = logging.getLogger(__name__)


async def detect_floorplan(repository: ListingRepository) -> None:
    """Detect square meters from floorplan images for all listings."""
    listings = await repository.get_listings()
    semaphore = asyncio.Semaphore(MAX_OCR_WORKERS)

    logger.info("Detecting floorplans for %d listings", len(listings))
    updated_listings = [
        listing
        for listing in await asyncio.gather(
            *[_calculate_sqm_ocr(listing, semaphore) for listing in listings]
        )
        if listing is not None
    ]
    logger.info("Finished floorplan detection, %d listings updated", len(updated_listings))
    await repository.upsert_listings(updated_listings)


async def _calculate_sqm_ocr(
    listing: Listing, semaphore: asyncio.Semaphore
) -> Listing | None:
    """Calculate square meters from floorplan images using OCR."""
    if listing.square_meters is not None:
        return None
    if not listing.floorplan_image_paths:
        listing.square_meters = 0
        return listing
    sqms: list[float] = []
    for floorplan_path in listing.floorplan_image_paths:
        async with semaphore:
            estimated_sqm, _ = await asyncio.to_thread(
                floorplan.calculate_ocr, floorplan_path
            )
            if estimated_sqm is not None:
                sqms.append(estimated_sqm)
    max_sqm = max(sqms, default=0)  # try once, if we fail, keep as 0
    listing.square_meters = max_sqm
    return listing