Upload 3 files

app.py

@@ -0,0 +1,83 @@
+import streamlit as st
+import torch
+import torchvision.transforms as tr
+from PIL import Image
+import numpy as np
+from model import CycleGAN, Discriminator, Generator
+
+# === Настройки ===
+mean = [0.5, 0.5, 0.5]
+std = [0.5, 0.5, 0.5]
+
+hyperparams = dict(
+    crop_size=256,
+)
+
+def get_transforms(mean, std,  crop_size=64):
+
+    val_transform = tr.Compose([
+        tr.Resize((crop_size, crop_size)),
+        tr.ToTensor(),
+        tr.Normalize(mean=mean, std=std)
+    ])
+
+    def de_normalize(tensor):
+        denorm = tr.Normalize(
+            mean=[-m / s for m, s in zip(mean, std)],
+            std=[1 / s for s in std]
+        )
+        return denorm(tensor.clone()).clamp(0, 1)
+
+    return val_transform, de_normalize
+
+
+val_transform, de_normalize = get_transforms(mean, std, **hyperparams)
+
+
+# === Загрузка модели ===
+@st.cache_resource
+def load_model():
+    checkpoint = torch.load("cycle_gan_face.pt", map_location="cpu", weights_only = False)
+    model = CycleGAN(Discriminator, Generator)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    return model
+
+
+model = load_model()
+
+# === Streamlit UI ===
+st.title("Обработка изображений через PyTorch модель")
+
+uploaded_file_1 = st.file_uploader("Загрузите изображение белого человека", type=["jpg", "jpeg", "png"], key="file1")
+uploaded_file_2 = st.file_uploader("Загрузите изображение черного человека", type=["jpg", "jpeg", "png"], key="file2")
+
+selected = st.radio("Выберите изображение для обработки", ["Первое", "Второе"])
+
+if uploaded_file_1 and uploaded_file_2:
+    image1 = Image.open(uploaded_file_1).convert("RGB")
+    image2 = Image.open(uploaded_file_2).convert("RGB")
+
+    if selected == "Первое":
+        selected_image = image1
+        st.image(selected_image, caption="Выбранное изображение", use_column_width=True)
+        tensor = val_transform(selected_image).unsqueeze(0)  # B x C x H x W
+
+        with torch.no_grad():
+            output = model.netG_A2B(tensor)
+    else:
+        selected_image = image2
+        st.image(selected_image, caption="Выбранное изображение", use_column_width=True)
+        tensor = val_transform(selected_image).unsqueeze(0)  # B x C x H x W
+
+        with torch.no_grad():
+            output = model.netG_B2A(tensor)
+
+
+    # Де-нормализуем и показываем результат
+    result_image = de_normalize(output.squeeze(0)).permute(1, 2, 0).numpy()
+    result_image = (result_image * 255).astype(np.uint8)
+
+    st.image(result_image, caption="Результат модели", use_column_width=True)
+else:
+    st.info("Пожалуйста, загрузите оба изображения.")

model.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Discriminator(nn.Module):
+    def __init__(self, dropout_prob=0.3):
+        super(Discriminator, self).__init__()
+
+        self.main = nn.Sequential(
+            nn.Conv2d(3, 64, 4, stride=2, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Dropout2d(p=dropout_prob),
+
+            nn.Conv2d(64, 128, 4, stride=2, padding=1),
+            nn.InstanceNorm2d(128),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Dropout2d(p=dropout_prob),
+
+            nn.Conv2d(128, 256, 4, stride=2, padding=1),
+            nn.InstanceNorm2d(256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Dropout2d(p=dropout_prob),
+
+            nn.Conv2d(256, 512, 4, padding=1),
+            nn.InstanceNorm2d(512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Dropout2d(p=dropout_prob),
+
+            nn.Conv2d(512, 1, 4, padding=1),
+        )
+
+    def forward(self, x):
+        x = self.main(x)
+        x = F.avg_pool2d(x, x.size()[2:])
+        x = torch.flatten(x, 1)
+        x = torch.sigmoid(x)
+        return x
+
+
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.main = nn.Sequential(
+            # Initial convolution block
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(3, 64, 7),
+            nn.InstanceNorm2d(64),
+            nn.ReLU(inplace=True),
+
+            # Downsampling
+            nn.Conv2d(64, 128, 3, stride=2, padding=1),
+            nn.InstanceNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 256, 3, stride=2, padding=1),
+            nn.InstanceNorm2d(256),
+            nn.ReLU(inplace=True),
+
+            # Residual blocks
+            ResidualBlock(256),
+            ResidualBlock(256),
+            ResidualBlock(256),
+            ResidualBlock(256),
+            ResidualBlock(256),
+            ResidualBlock(256),
+            ResidualBlock(256),
+            ResidualBlock(256),
+            ResidualBlock(256),
+
+            # Upsampling
+            nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1),
+            nn.InstanceNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
+            nn.InstanceNorm2d(64),
+            nn.ReLU(inplace=True),
+
+            # Output layer
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(64, 3, 7),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        return self.main(x)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels):
+        super(ResidualBlock, self).__init__()
+
+        self.res = nn.Sequential(nn.ReflectionPad2d(1),
+                                 nn.Conv2d(in_channels, in_channels, 3),
+                                 nn.InstanceNorm2d(in_channels),
+                                 nn.ReLU(inplace=True),
+                                 nn.ReflectionPad2d(1),
+                                 nn.Conv2d(in_channels, in_channels, 3),
+                                 nn.InstanceNorm2d(in_channels))
+
+    def forward(self, x):
+        return x + self.res(x)
+
+
+class CycleGAN(nn.Module):
+    def __init__(self, descriminator, generator):
+        super(CycleGAN, self).__init__()
+
+        self.netG_A2B = generator()
+        self.netG_B2A = generator()
+        self.netD_A   = descriminator()
+        self.netD_B   = descriminator()

requirements.txt

@@ -0,0 +1,3 @@
+torch==2.6.0
+torchvision==0.21.0
+streamlit==1.44.1