Skip to content

Commit

Permalink
fix(Proofs): Run price tag prediction only on PRICE_TAG proofs
Browse files Browse the repository at this point in the history
  • Loading branch information
raphodn committed Jan 15, 2025
1 parent 9268606 commit bd36b23
Showing 1 changed file with 8 additions and 9 deletions.
17 changes: 8 additions & 9 deletions open_prices/proofs/ml.py
Original file line number Diff line number Diff line change
Expand Up @@ -553,6 +553,9 @@ def run_and_save_price_tag_detection(
:return: the ProofPrediction instance created, or None if the prediction
already exists and overwrite is False
"""
if proof.type != proof_constants.TYPE_PRICE_TAG:
logger.debug("Skipping proof %s, not of type PRICE_TAG", proof.id)
return None

proof_prediction = ProofPrediction.objects.filter(
proof=proof, model_name=PRICE_TAG_DETECTOR_MODEL_NAME
Expand All @@ -570,10 +573,7 @@ def run_and_save_price_tag_detection(
proof.id,
PRICE_TAG_DETECTOR_MODEL_NAME,
)
if (
proof.type == proof_constants.TYPE_PRICE_TAG
and not PriceTag.objects.filter(proof=proof).exists()
):
if not PriceTag.objects.filter(proof=proof).exists():
logger.debug(
"Creating price tags from existing prediction for proof %s",
proof.id,
Expand All @@ -599,10 +599,9 @@ def run_and_save_price_tag_detection(
value=None,
max_confidence=max_confidence,
)
if proof.type == proof_constants.TYPE_PRICE_TAG:
create_price_tags_from_proof_prediction(
proof, proof_prediction, run_extraction=run_extraction
)
create_price_tags_from_proof_prediction(
proof, proof_prediction, run_extraction=run_extraction
)
return proof_prediction


Expand Down Expand Up @@ -663,7 +662,7 @@ def run_and_save_proof_prediction(
Currently, the following models are run:
- proof type classification model
- price tag detection model (objecct detector)
- price tag detection model (object detector)
:param proof_id: the ID of the proof to be classified
:param run_price_tag_extraction: whether to run the price tag extraction
Expand Down

0 comments on commit bd36b23

Please sign in to comment.