Create style_transfer.py

style_transfer.py

@@ -0,0 +1,20 @@
+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")