Views
No views yet
ComplexUNet architecture, a variant of the standard U-Net adapted for 256x256 images. It features a deep structure with 5 downsampling/upsampling stages and uses residual blocks for more stable training.torch and torchvision installed.model.py file in your project directory.inpainting_model_coco.pth file from the 'Files and versions' tab.1import torch
2from model import ComplexUNet # Import the class from model.py
3from PIL import Image
4import torchvision.transforms as T
5
6# --- Setup ---
7DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8MODEL_PATH = "inpainting_model_coco.pth" # <-- Make sure you've downloaded this file
9
10# --- Load Model ---
11# Note: Use base_channels=64 as it was during training
12model = ComplexUNet(base_channels=64)
13model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
14model.to(DEVICE)
15model.eval()
16
17print("Model loaded successfully!")
18
19# --- Example: Inpaint an image ---
20# 1. Load your masked image
21# masked_image = Image.open("path/to/your/masked_image.png").convert("RGB")
22#
23# 2. Create a tensor from your image
24# transform = T.Compose([
25# T.Resize(256),
26# T.CenterCrop(256),
27# T.ToTensor()
28# ])
29# masked_tensor = transform(masked_image).unsqueeze(0).to(DEVICE)
30#
31# 3. Get the reconstructed image
32# with torch.no_grad():
33# reconstructed_tensor = model(masked_tensor)
34#
35# 4. Convert tensor back to PIL Image
36# reconstructed_image = T.ToPILImage()(reconstructed_tensor.squeeze(0).cpu())
37# reconstructed_image.save("reconstructed_result.png")
38# print("Inpainting complete. Saved to reconstructed_result.png")