diff --git a/egomimic/models/stems/e1_image.py b/egomimic/models/stems/e1_image.py index 22a00768a..c935ee665 100644 --- a/egomimic/models/stems/e1_image.py +++ b/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 diff --git a/egomimic/models/stems/resnet_mlp.py b/egomimic/models/stems/resnet_mlp.py new file mode 100644 index 000000000..955e77f3e --- /dev/null +++ b/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)))