fix

app.py

@@ -3,7 +3,7 @@ import torch
 from torchvision import transforms as tr
 from PIL import Image
 from huggingface_hub import hf_hub_download
-from cycle_gan import CycleGAN, create_model_and_optimizer
+from cycle_gan import CycleGAN, create_model_and_optimizer, val_transform_a, val_transform_b, de_normalize_a, de_normalize_b
 
 @st.cache_resource  # кэширование
 def load_model():
@@ -30,16 +30,23 @@ uploaded_file = st.file_uploader(f"Upload your {query}", type=["png", "jpg", "jp
 if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
     image_size = image.size
-    image = image.resize((256,256))
-    image = tr.ToTensor()(image)
-    image = image.to(dtype = torch.float, device='cpu')
+    # image = image.resize((256,256))
+    image = tr.CenterCrop(min(image.size))(image)
+
     model.eval()
     with torch.no_grad():
         if style == 'ColoredToSketch':
-            generation = model.GA(image.view(1, *image.shape)).squeeze()
+            image = val_transform_a(image)
+            image = image.to(dtype = torch.float, device='cpu')
+            generation = model.GA(image.view(1, *image.shape)).detach().squeeze()
+            generation = de_normalize_b(generation)
         else:
-            generation = model.GB(image.view(1, *image.shape)).squeeze()
-        generation = tr.ToPILImage()(generation)
+            st.markdown(str(image.shape))
+            image = val_transform_b(image)
+            image = image.to(dtype = torch.float, device='cpu')
+            print(val_transform_b)
+            generation = model.GB(image.unsqueeze(0)).squeeze().detach()
+            generation = de_normalize_a(generation)
         resized = generation.resize(image_size, Image.Resampling.LANCZOS)
         st.image(resized, caption="Your result!", use_column_width=True)
 

cycle_gan.py

@@ -1,9 +1,54 @@
 
 import torch
 from torch import nn
+from torchvision import transforms as tr
+import numpy as np
 import functools
 import itertools
 
+channel_mean_a, channel_std_a = np.array([0.87415955, 0.84342639 ,0.8385736 ]), np.array([0.21790985, 0.24519696, 0.24330734])
+channel_mean_b, channel_std_b = np.array([0.95326045, 0.95326045, 0.95326045]), np.array([0.13081782, 0.13081782, 0.13081782])
+
+def get_transforms(mean, std, crop_size):
+    train_transform = tr.Compose([
+        tr.ToPILImage(),
+        tr.RandomCrop(crop_size),
+        tr.RandomHorizontalFlip(0.5),
+        tr.RandomVerticalFlip(0.5),
+        # tr.ColorJitter(brightness=0.05, contrast=0.1, saturation=0.1, hue=0.1),
+        # tr.RandomRotation(30),
+        # tr.RandomVerticalFlip(0.5),
+        tr.Resize(256),
+        tr.ToTensor(),
+        tr.Normalize(mean, std)
+    ])
+
+    val_transform = tr.Compose([
+        tr.ToPILImage(),
+        tr.Resize(256),
+        tr.ToTensor(),
+        tr.Normalize(mean, std)
+    ])
+
+    def de_normalize(img, normalized=True):
+        # print(type(img))
+        res = img.cpu()*std.reshape(-1,1,1) + mean.reshape(-1,1,1)
+        return tr.ToPILImage()(res)
+        # img = img.detach().cpu().numpy().transpose((1, 2, 0))
+        # return img
+        # return img * std + mean
+
+    return train_transform, val_transform, de_normalize
+
+
+# Ваши гиперпараметры
+hyperparams = dict(
+    [('crop_size',224)]
+)
+# transform-ы для A и B
+train_transform_a, val_transform_a, de_normalize_a = get_transforms(channel_mean_a, channel_std_a, **hyperparams)
+train_transform_b, val_transform_b, de_normalize_b = get_transforms(channel_mean_b, channel_std_b, **hyperparams)
+
 class Discriminator(nn.Module):
   def __init__(self, c_in, c_out_init=64, n_layers=3, norm_layer=nn.BatchNorm2d):
     super().__init__()