Merge pull request #208 from stergioc/hoptimus

Addin H-optimus-0 Model

Author: quentinblampey

sopa/patches/models/__init__.py

@@ -2,12 +2,14 @@
 from .histo_ssl import HistoSSLFeatures
 from .dinov2 import DINOv2Features
 from .dummy import DummyFeatures
+from .hoptimus0 import HOPTIMUSFeatures
 
-__all__ = ["Resnet50Features", "HistoSSLFeatures", "DINOv2Features", "DummyFeatures"]
+__all__ = ["Resnet50Features", "HistoSSLFeatures", "DINOv2Features", "DummyFeatures", "HOPTIMUSFeatures"]
 
 available_models = {
     "resnet50": Resnet50Features,
     "histo_ssl": HistoSSLFeatures,
     "dinov2": DINOv2Features,
     "dummy": DummyFeatures,
+    "hoptimus0": HOPTIMUSFeatures,
 }

sopa/patches/models/hoptimus0.py

@@ -0,0 +1,20 @@
+from torch import nn
+from torchvision import transforms
+
+
+class HOPTIMUSFeatures(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        import timm
+
+        self.model = timm.create_model("hf_hub:bioptimus/H-optimus-0", pretrained=True)
+
+    def forward(self, x):
+        transform = transforms.Compose(
+            [
+                transforms.Normalize(mean=(0.707223, 0.578729, 0.703617), std=(0.211883, 0.230117, 0.177517)),
+            ]
+        )
+
+        return self.model(transform(x))