add: judge configs to the scores attribute in SpatialRelationshipMetr…

…ic2D to be logged to the wandb table
wandb · Jun 21, 2024 · 59b166e · 59b166e
1 parent 6e46579
commit 59b166e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/hemm/metrics/spatial_relationship/judges/detr.py b/hemm/metrics/spatial_relationship/judges/detr.py
@@ -18,10 +18,12 @@ class DETRSpatialRelationShipJudge(weave.Model):
     Args:
         model_address (str, optional): The address of the model to use.
         revision (str, optional): The revision of the model to use.
+        name (str, optional): The name of the judge model
     """
 
     model_address: str = "facebook/detr-resnet-50"
     revision: str = "no_timm"
+    name: str = "detr_spatial_relationship_judge"
     _feature_extractor: DetrImageProcessor = None
     _object_detection_model: DetrForObjectDetection = None
 

diff --git a/hemm/metrics/spatial_relationship/spatial_relationship_2d.py b/hemm/metrics/spatial_relationship/spatial_relationship_2d.py
@@ -62,6 +62,7 @@ def __init__(
     ) -> None:
         self.judge = judge
         self.judge._initialize_models()
+        self.judge_config = self.judge.model_dump(mode="json")
         self.iou_threshold = iou_threshold
         self.distance_threshold = distance_threshold
         self.name = name
@@ -193,7 +194,8 @@ def compose_judgement(
                         base64_decode_image(annotated_image)
                         if isinstance(annotated_image, str)
                         else annotated_image
-                    )
+                    ),
+                    "judge_config": self.judge_config,
                 },
             }
         )
@@ -206,6 +208,7 @@ def compose_judgement(
                     else annotated_image
                 )
             },
+            "judge_config": self.judge_config,
         }
 
     @weave.op()