[Pipeline 12/12] replace campaign-named image stem with shared ResNet-MLP implementation

egomimic/models/stems/e1_image.py

@@ -1,27 +1,5 @@
-"""E1 image path: ImageNet normalization, ResNet, then the HPT MLP stem.
+"""Historical import name for the shared ResNet-MLP stem; removed with old callers."""
 
-The campaign's encoder and stem were separate in the previous HPT runtime.
-Keeping both here preserves that path while HPTStemStage owns latent pooling.
-"""
+from egomimic.models.stems.resnet_mlp import ResNetMLPImageStem
 
-from torchvision import transforms
-
-from egomimic.models.stems.hpt_stems import MLPPolicyStem, PolicyStem, ResNet
-
-
-class E1ImageStem(PolicyStem):
-    def __init__(self, output_dim=256, weights="DEFAULT", **kwargs):
-        super().__init__(**kwargs)
-        self.encoder = ResNet(output_dim=output_dim, weights=weights)
-        self.mlp = MLPPolicyStem(
-            input_dim=output_dim, output_dim=output_dim, widths=[output_dim]
-        )
-        self.jitter = transforms.ColorJitter(0.1, 0.1, 0.1, 0.05)
-        self.normalize = transforms.Normalize(
-            [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
-        )
-
-    def forward(self, images):
-        if self.training:
-            images = self.jitter(images)
-        return self.mlp(self.encoder(self.normalize(images)))
+E1ImageStem = ResNetMLPImageStem

egomimic/models/stems/resnet_mlp.py

@@ -0,0 +1,27 @@
+"""ImageNet-normalized ResNet features followed by an HPT MLP stem.
+
+Color jitter is applied only during training. Spatial tokens are preserved
+for latent pooling in HPTStemStage.
+"""
+
+from torchvision import transforms
+
+from egomimic.models.stems.hpt_stems import MLPPolicyStem, PolicyStem, ResNet
+
+
+class ResNetMLPImageStem(PolicyStem):
+    def __init__(self, output_dim=256, weights="DEFAULT", **kwargs):
+        super().__init__(**kwargs)
+        self.encoder = ResNet(output_dim=output_dim, weights=weights)
+        self.mlp = MLPPolicyStem(
+            input_dim=output_dim, output_dim=output_dim, widths=[output_dim]
+        )
+        self.jitter = transforms.ColorJitter(0.1, 0.1, 0.1, 0.05)
+        self.normalize = transforms.Normalize(
+            [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
+        )
+
+    def forward(self, images):
+        if self.training:
+            images = self.jitter(images)
+        return self.mlp(self.encoder(self.normalize(images)))