| def transfer_style(image_path, target_style_name): | |
| G = ConditionalGenerator(num_styles=3) | |
| G.load_state_dict(torch.load("generator.pth")) | |
| img = Image.open(image_path).convert("RGB") | |
| transform = T.Compose([T.Resize(256), T.ToTensor(), T.Normalize(0.5, 0.5)]) | |
| img_tensor = transform(img).unsqueeze(0) | |
| style_id = style_name_to_id[target_style_name] | |
| with torch.no_grad(): | |
| stylized = G(img_tensor, torch.tensor([style_id])) | |
| result = ToPILImage()((stylized.squeeze() + 1) / 2) # Денормализация [0, 1] | |
| result.save("stylized.jpg") |