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)