Create model.py

model.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super(UNet, self).__init__()
+        # Definição simplificada da UNet
+        self.enc1 = nn.Sequential(
+            nn.Conv2d(in_channels, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1),
+            nn.ReLU()
+        )
+        self.pool = nn.MaxPool2d(2)
+        self.dec1 = nn.Sequential(
+            nn.ConvTranspose2d(64, 64, kernel_size=2, stride=2),
+            nn.ReLU(),
+            nn.Conv2d(64, out_channels, kernel_size=1)
+        )
+
+    def forward(self, x):
+        x1 = self.enc1(x)
+        x2 = self.pool(x1)
+        x3 = self.dec1(x2)
+        return x3