Update app.py

app.py

@@ -9,6 +9,13 @@ import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
 import os
+import sys
+
+# Проверка правильного запуска
+if not 'streamlit' in sys.modules:
+    print("Пожалуйста, запустите приложение с помощью команды:")
+    print("streamlit run dcgan_app.py")
+    sys.exit(1)
 
 # Параметры
 nc = 3  # Количество каналов в изображении
@@ -68,14 +75,16 @@ class Discriminator(nn.Module):
     def forward(self, input):
         return self.main(input)
 
-# Настройка Streamlit
-st.title('DCGAN Training and Generation')
-
-# Создание боковой панели
-st.sidebar.title('Параметры')
-num_epochs = st.sidebar.slider('Количество эпох', 1, 50, 5)
-batch_size = st.sidebar.slider('Размер батча', 16, 128, 64)
-lr = st.sidebar.number_input('Скорость обучения', 0.0001, 0.01, 0.0002)
+# Функция для обработки тензора в изображение
+def process_tensor_to_image(tensor):
+    try:
+        img = vutils.make_grid(tensor, padding=2, normalize=True)
+        img = np.transpose(img, (1, 2, 0))
+        img = ((img + 1) * 127.5).astype(np.uint8)
+        return Image.fromarray(img)
+    except Exception as e:
+        st.error(f"Ошибка при обработке изображения: {e}")
+        return None
 
 # Загрузка данных
 @st.cache_data
@@ -106,9 +115,7 @@ def generate_images(netG, num_images=64):
     with torch.no_grad():
         noise = torch.randn(num_images, nz, 1, 1, device=device)
         fake = netG(noise).detach().cpu()
-    img = vutils.make_grid(fake, padding=2, normalize=True)
-    img = np.transpose(img, (1, 2, 0))
-    return img
+    return process_tensor_to_image(fake)
 
 # Функция обучения
 def train_model():
@@ -158,7 +165,7 @@ def train_model():
             errD = errD_real + errD_fake
             optimizerD.step()
 
-############################
+            ############################
             # (2) Обновление генератора
             ###########################
             netG.zero_grad()
@@ -178,11 +185,9 @@ def train_model():
                 
                 # Показать промежуточные результаты
                 if i % 500 == 0:
-                    with torch.no_grad():
-                        fake = netG(torch.randn(64, nz, 1, 1, device=device)).detach().cpu()
-                    img = vutils.make_grid(fake, padding=2, normalize=True)
-                    img = np.transpose(img, (1, 2, 0))
-                    st.image(img, caption=f'Эпоха {epoch}, Batch {i}')
+                    img = process_tensor_to_image(fake)
+                    if img is not None:
+                        st.image(img, caption=f'Эпоха {epoch}, Batch {i}')
         
         # Обновление прогресс бара
         progress_bar.progress((epoch + 1) / num_epochs)
@@ -193,7 +198,14 @@ def train_model():
 
 # Основной интерфейс Streamlit
 def main():
-    st.sidebar.title('DCGAN Control Panel')
+    st.title('DCGAN Training and Generation')
+    
+    # Настройка боковой панели
+    st.sidebar.title('Параметры')
+    global num_epochs, batch_size, lr
+    num_epochs = st.sidebar.slider('Количество эпох', 1, 50, 5)
+    batch_size = st.sidebar.slider('Размер батча', 16, 128, 64)
+    lr = st.sidebar.number_input('Скорость обучения', 0.0001, 0.01, 0.0002)
     
     # Выбор режима
     mode = st.sidebar.selectbox('Выберите режим', 
@@ -212,7 +224,8 @@ def main():
             # Генерация финальных изображений
             st.subheader('Финальные сгенерированные изображения')
             final_images = generate_images(netG)
-            st.image(final_images, caption='Финальные сгенерированные изображения')
+            if final_images is not None:
+                st.image(final_images, caption='Финальные сгенерированные изображения')
     
     elif mode == 'Генерация':
         if os.path.exists('generator.pth'):
@@ -224,20 +237,16 @@ def main():
             
             if st.button('Сгенерировать изображения'):
                 images = generate_images(netG, num_images)
-                st.image(images, caption='Сгенерированные изображения')
-                
-                # Опция сохранения
-                if st.button('Сохранить изображения'):
-                    im = Image.fromarray((images * 255).astype(np.uint8))
-                    im.save('generated_images.png')
-                    st.success('Изображения сохранены!')
+                if images is not None:
+                    st.image(images, caption='Сгенерированные изображения')
+                    
+                    # Опция сохранения
+                    if st.button('Сохранить изображения'):
+                        images.save('generated_images.png')
+                        st.success('Изображения сохранены!')
         else:
             st.error('Модель не найдена. Пожалуйста, сначала обучите модель.')
 
-# Запуск приложения
-if __name__ == '__main__':
-    main()
-
 # Дополнительные настройки
 st.sidebar.markdown("""
 ## О проекте
@@ -253,11 +262,9 @@ st.sidebar.markdown("""
 
 # Настройки кэширования
 if st.sidebar.checkbox('Очистить кэш'):
-    st.caching.clear_cache()
+    st.cache_data.clear()
     st.success('Кэш очищен!')
 
-# Дополнительные метрики
-if st.sidebar.checkbox('Показать дополнительные метрики'):
-    st.sidebar.write(f'Размер батча: {batch_size}')
-    st.sidebar.write(f'Количество эпох: {num_epochs}')
-    st.sidebar.write(f'Скорость обучения: {lr}')
+# Запуск приложения
+if __name__ == '__main__':
+    main()