From 8de66d87f1671cae1a24b393975550242217c184 Mon Sep 17 00:00:00 2001
From: unanmed <1319491857@qq.com>
Date: Tue, 31 Mar 2026 22:24:25 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E9=9A=8F=E6=9C=BA=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=A2=9E=E5=BC=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ginka/dataset.py         | 20 ++++++++++++++++++--
 ginka/maskGIT/maskGIT.py | 23 +++++++++++++++++++++++
 ginka/maskGIT/model.py   | 19 +++++--------------
 3 files changed, 46 insertions(+), 16 deletions(-)
 create mode 100644 ginka/maskGIT/maskGIT.py

diff --git a/ginka/dataset.py b/ginka/dataset.py
index 2450b6d..5558942 100644
--- a/ginka/dataset.py
+++ b/ginka/dataset.py
@@ -28,10 +28,26 @@ class GinkaMaskGITDataset(Dataset):
     def __getitem__(self, idx):
         item = self.data[idx]
         
-        target = torch.LongTensor(item['map']) # [H, W]
-        cond = torch.FloatTensor(item['val']) # [cond_dim]
+        target_np = np.array(item['map'])
         heatmap = np.array(item['heatmap'], dtype=np.float32)
         
+        # 数据增强
+        if np.random.rand() > 0.5:
+            k = np.random.randint(0, 4)
+            target_np = np.rot90(target_np, k)
+            heatmap = np.rot90(heatmap, k)
+
+        if np.random.rand() > 0.5:
+            target_np = np.fliplr(target_np)
+            heatmap = np.fliplr(heatmap)
+
+        if np.random.rand() > 0.5:
+            target_np = np.flipud(target_np)
+            heatmap = np.flipud(heatmap)
+        
+        target = torch.LongTensor(target_np) # [H, W]
+        cond = torch.FloatTensor(item['val']) # [cond_dim]
+        
         if random.random() < 0.5:
             size = random.randint(self.blur_min, self.blur_max)
             if size % 2 == 0:
diff --git a/ginka/maskGIT/maskGIT.py b/ginka/maskGIT/maskGIT.py
new file mode 100644
index 0000000..15d0d69
--- /dev/null
+++ b/ginka/maskGIT/maskGIT.py
@@ -0,0 +1,23 @@
+import torch.nn as nn
+
+class MaskGIT(nn.Module):
+    def __init__(
+        self, d_model=256, dim_ff=512, nhead=8, num_layers=4,
+    ):
+        super().__init__()
+        self.encoder = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_ff, batch_first=True, activation='gelu'),
+            num_layers=num_layers
+        )
+
+        self.decoder = nn.TransformerDecoder(
+            nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_ff, batch_first=True, activation='gelu'),
+            num_layers=num_layers
+        )
+        
+    def forward(self, x):
+        # x: [B, L, d_model]
+        m = self.encoder(x)
+        out = self.decoder(x, m)
+        return out
+        
\ No newline at end of file
diff --git a/ginka/maskGIT/model.py b/ginka/maskGIT/model.py
index 7f86800..2ea5171 100644
--- a/ginka/maskGIT/model.py
+++ b/ginka/maskGIT/model.py
@@ -3,6 +3,7 @@ import torch
 import torch.nn as nn
 from ..utils import print_memory
 from .cond import GinkaMaskGITCond
+from .maskGIT import MaskGIT
 
 class GinkaMaskGIT(nn.Module):
     def __init__(
@@ -16,14 +17,7 @@ class GinkaMaskGIT(nn.Module):
         
         self.cond_encoder = GinkaMaskGITCond(cond_dim=cond_dim, heatmap_channel=heatmap_channel, output_dim=d_model)
         
-        self.encoder = nn.TransformerEncoder(
-            nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_ff, batch_first=True, activation='gelu'),
-            num_layers=num_layers
-        )
-        self.decoder = nn.TransformerDecoder(
-            nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_ff, batch_first=True, activation='gelu'),
-            num_layers=num_layers
-        )
+        self.transformer = MaskGIT(d_model=d_model, dim_ff=dim_ff, nhead=nhead, num_layers=num_layers)
         
         self.output_fc = nn.Sequential(
             nn.Linear(d_model, num_classes)
@@ -43,11 +37,9 @@ class GinkaMaskGIT(nn.Module):
         heatmap = heatmap.view(B, C, H * W).permute(0, 2, 1)
         x = self.tile_embedding(map) + heatmap
         x = torch.cat([cond.unsqueeze(1), x], dim=1) + self.pos_embedding
+        x = self.transformer(x)
         
-        m = self.encoder(x)
-        out = self.decoder(x, m)
-        
-        logits = self.output_fc(out)
+        logits = self.output_fc(x)
         
         return logits[:, :-1, :]
         
@@ -74,7 +66,6 @@ if __name__ == "__main__":
     print(f"输出形状: output={output.shape}")
     print(f"Tile Embedding parameters: {sum(p.numel() for p in model.tile_embedding.parameters())}")
     print(f"Condition Encoder parameters: {sum(p.numel() for p in model.cond_encoder.parameters())}")
-    print(f"Encoder parameters: {sum(p.numel() for p in model.encoder.parameters())}")
-    print(f"Decoder parameters: {sum(p.numel() for p in model.decoder.parameters())}")
+    print(f"MaskGIT parameters: {sum(p.numel() for p in model.transformer.parameters())}")
     print(f"Output parameters: {sum(p.numel() for p in model.output_fc.parameters())}")
     print(f"Total parameters: {sum(p.numel() for p in model.parameters())}")