utils/worker.py

import hashlib
import json
import os

from src.detection import ProductDetector
from src.grouping import ProductGrouper
from src.visualize import visualize

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODEL_PATH = os.path.join(BASE_DIR, "models", "best.pt")
DATA_DIR = os.getenv("APP_DATA_DIR", BASE_DIR)

detector = ProductDetector(MODEL_PATH)
grouper = ProductGrouper()

OUTPUT_FOLDER = os.path.join(DATA_DIR, "outputs")
RESULT_FOLDER = os.path.join(DATA_DIR, "results")

os.makedirs(OUTPUT_FOLDER, exist_ok=True)
os.makedirs(RESULT_FOLDER, exist_ok=True)


def _sha256(path):
    if not os.path.exists(path):
        return None
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


RUNTIME_META = {
    "model_path": MODEL_PATH,
    "model_sha256": _sha256(MODEL_PATH),
    "yolo_device": str(detector.device),
    "clip_device": str(grouper.device),
    "inference_seed": int(grouper.seed),
    "detector_conf": float(detector.conf_threshold),
    "detector_iou": float(detector.iou_threshold),
    "detector_imgsz": int(detector.imgsz),
    "similarity_threshold": float(grouper.similarity_threshold),
    "similarity_margin": float(grouper.similarity_margin),
    "similarity_round_decimals": int(grouper.similarity_round_decimals),
    "row_vertical_threshold": int(grouper.row_vertical_threshold),
    "padding": int(grouper.padding),
}

def process_image(image_path, image_id):

    detections = detector.detect(image_path)

    grouped = grouper.group_products(image_path, detections)

    output_path = os.path.join(OUTPUT_FOLDER, f"{image_id}.jpg")

    visualize(image_path, grouped, output_path)

    result = {
        "image_id": image_id,
        "total_products": len(grouped),
        "detections": grouped,
        "visualization": output_path,
        "meta": RUNTIME_META,
    }

    with open(os.path.join(RESULT_FOLDER, f"{image_id}.json"), "w") as f:
        json.dump(result, f)

    return result