Refine DRF preprocessing and body-prior pipeline

opengait/modeling/models/drf.py

@@ -8,7 +8,7 @@ from jaxtyping import Float, Int
 
 from einops import rearrange
 
-from ..base_model import BaseModel
+from ..base_model_body import BaseModelBody
 from ..modules import (
     HorizontalPoolingPyramid,
     PackSequenceWrapper,
@@ -18,7 +18,7 @@ from ..modules import (
 )
 
 
-class DRF(BaseModel):
+class DRF(BaseModelBody):
     """Dual Representation Framework from arXiv:2509.00872v1."""
 
     def build_network(self, model_cfg: dict[str, Any]) -> None:
@@ -43,9 +43,10 @@ class DRF(BaseModel):
             list[str],
             list[str],
             Int[torch.Tensor, "1 batch"] | None,
+            Float[torch.Tensor, "batch pairs metrics"],
         ],
     ) -> dict[str, dict[str, Any]]:
-        ipts, pids, labels, _, seqL = inputs
+        ipts, pids, labels, _, seqL, key_features = inputs
         label_ids = torch.as_tensor(
             [LABEL_MAP[str(label).lower()] for label in labels],
             device=pids.device,
@@ -58,15 +59,12 @@ class DRF(BaseModel):
         else:
             heatmaps = rearrange(heatmaps, "n s c h w -> n c s h w")
 
-        pav_seq = ipts[1]
-        pav = aggregate_sequence_features(pav_seq, seqL)
-
         outs = self.Backbone(heatmaps)
         outs = self.TP(outs, seqL, options={"dim": 2})[0]
 
         feat = self.HPP(outs)
         embed_1 = self.FCs(feat)
-        embed_1 = self.PGA(embed_1, pav)
+        embed_1 = self.PGA(embed_1, key_features)
 
         embed_2, logits = self.BNNecks(embed_1)
         del embed_2
@@ -120,24 +118,6 @@ class PAVGuidedAttention(nn.Module):
         return embeddings * channel_att * spatial_att
 
 
-def aggregate_sequence_features(
-    sequence_features: Float[torch.Tensor, "batch seq pairs metrics"],
-    seqL: Int[torch.Tensor, "1 batch"] | None,
-) -> Float[torch.Tensor, "batch pairs metrics"]:
-    if seqL is None:
-        return sequence_features.mean(dim=1)
-
-    lengths = seqL[0].tolist()
-    flattened = sequence_features.squeeze(0)
-    aggregated = []
-    start = 0
-    for length in lengths:
-        end = start + int(length)
-        aggregated.append(flattened[start:end].mean(dim=0))
-        start = end
-    return torch.stack(aggregated, dim=0)
-
-
 LABEL_MAP: dict[str, int] = {
     "negative": 0,
     "neutral": 1,