Update app.py

app.py

@@ -10,6 +10,31 @@ from scipy.interpolate import griddata
 import matplotlib.pyplot as plt
 from utils import azi_diff
 
+class AttentionBlock(nn.Module):
+    def __init__(self, input_dim, num_heads, ff_dim, rate=0.2):
+        super(AttentionBlock, self).__init__()
+        self.attention = nn.MultiheadAttention(embed_dim=input_dim, num_heads=num_heads)
+        self.dropout1 = nn.Dropout(rate)
+        self.layer_norm1 = nn.LayerNorm(input_dim)
+
+        self.ffn = nn.Sequential(
+            nn.Linear(input_dim, ff_dim),
+            nn.ReLU(),
+            nn.Dropout(rate),
+            nn.Linear(ff_dim, input_dim),
+            nn.Dropout(rate)
+        )
+        self.layer_norm2 = nn.LayerNorm(input_dim)
+
+    def forward(self, x):
+        attn_output, _ = self.attention(x, x, x)
+        attn_output = self.dropout1(attn_output)
+        out1 = self.layer_norm1(attn_output + x)
+
+        ffn_output = self.ffn(out1)
+        out2 = self.layer_norm2(ffn_output + out1)
+        return out2
+
 class TextureContrastClassifier(nn.Module):
     def __init__(self, input_shape, num_heads=4, key_dim=64, ff_dim=256, rate=0.1):
         super(TextureContrastClassifier, self).__init__()