Upload 4 files

main.py

@@ -0,0 +1,34 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+from watermark_remover import WatermarkRemover
+import numpy as np
+
+image_path = "path to your test image"  # Replace with the path to your test image
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load the trained model
+model = WatermarkRemover().to(device)
+model_path = "path to your model.pth"  # Replace with the path to your saved model
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
+
+transform = transforms.Compose([transforms.Resize((256, 256)),
+                                transforms.ToTensor(),])
+watermarked_image = Image.open(image_path).convert("RGB")
+original_size = watermarked_image.size  
+input_tensor = transform(watermarked_image).unsqueeze(0).to(device)
+
+with torch.no_grad():
+    output_tensor = model(input_tensor)
+
+predicted_image = output_tensor.squeeze(0).cpu().permute(1, 2, 0).clamp(0, 1).numpy()
+predicted_pil = Image.fromarray((predicted_image * 255).astype(np.uint8))
+predicted_pil = predicted_pil.resize(original_size, Image.Resampling.LANCZOS)
+predicted_pil.save("predicted_image.jpg", quality=100)
+
+
+
+
+
+

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfceb0199f80877802c6f0c5abf412f39ee82a634cbbcf072d05ce7c06e33d3b
+size 125529866

watermark_remover.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+
+class WatermarkRemover(nn.Module):
+    def __init__(self):
+        super(WatermarkRemover, self).__init__()
+        self.enc1 = self.conv_block(3, 64)
+        self.enc2 = self.conv_block(64, 128)
+        self.enc3 = self.conv_block(128, 256)
+        self.enc4 = self.conv_block(256, 512)
+
+        self.bottleneck = self.conv_block(512, 1024)
+
+        self.dec4 = self.conv_block(1024 + 512, 512)
+        self.dec3 = self.conv_block(512 + 256, 256)
+        self.dec2 = self.conv_block(256 + 128, 128)
+        self.dec1 = self.conv_block(128 + 64, 64)
+
+        self.final_layer = nn.Conv2d(64, 3, kernel_size=1)
+
+    def conv_block(self, in_channels, out_channels):
+        return nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(nn.MaxPool2d(2)(e1))
+        e3 = self.enc3(nn.MaxPool2d(2)(e2))
+        e4 = self.enc4(nn.MaxPool2d(2)(e3))
+
+        b = self.bottleneck(nn.MaxPool2d(2)(e4))
+
+        d4 = self.dec4(torch.cat((nn.Upsample(scale_factor=2)(b), e4), dim=1))
+        d3 = self.dec3(torch.cat((nn.Upsample(scale_factor=2)(d4), e3), dim=1))
+        d2 = self.dec2(torch.cat((nn.Upsample(scale_factor=2)(d3), e2), dim=1))
+        d1 = self.dec1(torch.cat((nn.Upsample(scale_factor=2)(d2), e1), dim=1))
+
+        return self.final_layer(d1)
+