upload model to huggingface

model.py

@@ -0,0 +1,97 @@
+import torch
+import torchvision
+import torch.nn as nn
+import torch.optim as optim
+import math
+
+class ImprovedEfficientBackbone(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.efficientnet = torchvision.models.efficientnet_b0(weights=torchvision.models.EfficientNet_B0_Weights.IMAGENET1K_V1)
+        self.features = self.efficientnet.features
+
+    def forward(self, x):
+        return self.features(x) 
+
+
+class ImprovedPatchEmbedding(nn.Module):
+    def __init__(self, in_channels=1280, embed_dim=384):
+        super().__init__()
+        self.proj = nn.Linear(in_channels, embed_dim)
+
+    def forward(self, x):
+        """
+        Input:  [B, 1280, 7, 7]
+        Output: [B, 49, 384]
+        """
+        B, C, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2) 
+        x = self.proj(x)  
+        return x
+
+
+class ImprovedViTBlock(nn.Module):
+    def __init__(self, embed_dim=384, num_heads=4, mlp_ratio=4):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim * mlp_ratio),
+            nn.GELU(),
+            nn.Linear(embed_dim * mlp_ratio, embed_dim)
+        )
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, x):
+        x = x + self.dropout(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0])
+        x = x + self.dropout(self.mlp(self.norm2(x)))
+        return x
+
+
+
+class ImprovedEfficientViT(nn.Module):
+    def __init__(self, embed_dim=384, depth=8, num_heads=4):
+        super().__init__()
+        self.backbone = ImprovedEfficientBackbone()
+        self.patch_embed = ImprovedPatchEmbedding(embed_dim=embed_dim)
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
+        self.register_buffer("pos_embed", self._get_sinusoidal_encoding(50, embed_dim))  # Use sin/cos
+
+        self.blocks = nn.ModuleList([ImprovedViTBlock(embed_dim, num_heads) for _ in range(depth)])
+
+        self.head = nn.Sequential(
+              nn.LayerNorm(embed_dim),
+              nn.Linear(embed_dim, 128),
+              nn.GELU(),
+              nn.Linear(128, 1)
+          )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+
+    def _get_sinusoidal_encoding(self, seq_len, dim):
+        pe = torch.zeros(seq_len, dim)
+        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe.unsqueeze(0)  
+
+    def forward(self, x):
+        features = self.backbone(x)                
+        tokens = self.patch_embed(features)         
+
+        B = tokens.shape[0]
+        cls_tokens = self.cls_token.expand(B, -1, -1) 
+        x = torch.cat((cls_tokens, tokens), dim=1)    
+        x = x + self.pos_embed[:, :x.size(1), :]      
+
+        for block in self.blocks:
+            x = block(x)
+
+        cls_final = x[:, 0]  
+        return self.head(cls_final) 

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+opencv-python
+mtcnn
+Pillow