Create modeling.py

modeling.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+import timm
+from huggingface_hub import PyTorchModelHubMixin
+
+class KeypointModel(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, heatmap_size=(512, 1024), **kwargs):
+        super().__init__()
+        config_heatmap_size = kwargs.get("config", {}).get("heatmap_size", heatmap_size)
+
+        backbone = timm.create_model('convnextv2_base.fcmae_ft_in22k_in1k_384', pretrained=False)
+
+        self.feature_extractor = nn.Sequential(*list(backbone.children())[:-2])
+        in_channels = backbone.num_features
+        self.head = nn.Sequential(
+            nn.Conv2d(in_channels, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Upsample(size=config_heatmap_size, mode='bilinear', align_corners=False),
+            nn.Conv2d(256, 1, kernel_size=1)
+        )
+
+    def forward(self, image):
+        features = self.feature_extractor(image)
+        heatmap = self.head(features)
+        return heatmap