predict.py

import os
import argparse
import torch
from PIL import Image
from torchvision import transforms
from model import UNet
from tqdm import tqdm

def predict(args):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Using device: {device}')

    # Load Model
    model = UNet(n_channels=3, n_classes=3, bilinear=True)
    
    if not os.path.exists(args.model_path):
        print(f"Error: Model not found at {args.model_path}")
        return

    model.load_state_dict(torch.load(args.model_path, map_location=device))
    model.to(device)
    model.eval()

    # Setup output directory
    os.makedirs(args.output_dir, exist_ok=True)

    # Transforms
    transform = transforms.Compose([
        transforms.Resize((args.img_size, args.img_size)),
        transforms.ToTensor()
    ])

    # Process images
    images = sorted([f for f in os.listdir(args.input_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
    
    print(f"Found {len(images)} images in {args.input_dir}")

    with torch.no_grad():
        for img_name in tqdm(images):
            img_path = os.path.join(args.input_dir, img_name)
            image = Image.open(img_path).convert("RGB")
            original_size = image.size
            
            # Preprocess
            input_tensor = transform(image).unsqueeze(0).to(device)

            # Inference
            output = model(input_tensor)

            # Postprocess
            output = output.squeeze(0).cpu()
            output = torch.clamp(output, 0, 1)
            output_img = transforms.ToPILImage()(output)
            
            # Resize back to original size
            output_img = output_img.resize(original_size, Image.BICUBIC)

            # Save
            save_path = os.path.join(args.output_dir, img_name)
            output_img.save(save_path)

    print(f"Done! Results saved to {args.output_dir}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Inference for Handwriting Removal')
    parser.add_argument('--input_dir', type=str, default='data/test', help='Path to input images')
    parser.add_argument('--output_dir', type=str, default='results', help='Path to save results')
    parser.add_argument('--model_path', type=str, default='checkpoints/best_model.pth', help='Path to trained model')
    parser.add_argument('--img_size', type=int, default=576, help='Image size used during training')
    
    args = parser.parse_args()
    predict(args)