| 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) | |