Update app.py

app.py

@@ -1,15 +1,9 @@
 import os
-import spaces
 import gradio as gr
 import numpy as np
-import tensorflow as tf
-from keras.applications.mobilenet_v2 import MobileNetV2, preprocess_input
-from keras.models import Model
-import matplotlib.pyplot as plt
-import logging
-from skimage.transform import resize
-from PIL import Image, ImageEnhance, ImageFilter
+from PIL import Image
 from tqdm import tqdm
+import matplotlib.pyplot as plt
 
 # Disable GPU usage by default
 os.environ['CUDA_VISIBLE_DEVICES'] = ''
@@ -42,7 +36,6 @@ class SwarmNeuralNetwork:
         self.image_shape = image_shape
         self.agents = [SwarmAgent(self.random_position(), self.random_velocity()) for _ in range(num_agents)]
         self.target_image = self.load_target_image(target_image_path)
-        self.mobilenet = self.load_mobilenet_model()
 
     def random_position(self):
         return np.random.randn(*self.image_shape)
@@ -51,13 +44,9 @@ class SwarmNeuralNetwork:
         return np.random.randn(*self.image_shape) * 0.01
 
     def load_target_image(self, img_path):
-        img = Image.open(img_path).resize((self.image_shape[1], self.image_shape[0]))
+        img = Image.open(img_path).convert('RGB').resize((self.image_shape[1], self.image_shape[0]))
         return np.array(img) / 127.5 - 1
 
-    def load_mobilenet_model(self):
-        mobilenet = MobileNetV2(weights='imagenet', include_top=False, input_shape=(128, 128, 3))
-        return Model(inputs=mobilenet.input, outputs=mobilenet.get_layer('block_13_expand_relu').output)
-
     def update_agents(self, timestep):
         for agent in self.agents:
             # Convert agent's position and target image into HDC space
@@ -82,23 +71,14 @@ class SwarmNeuralNetwork:
     def train(self, epochs):
         for epoch in tqdm(range(epochs), desc="Training"):
             self.update_agents(epoch)
-            generated_image = self.generate_image()
-
-            # Display the generated image
-            self.display_image(generated_image, title=f'Epoch {epoch}')
 
-    def display_image(self, image, title=''):
-        plt.imshow(image)
-        plt.title(title)
-        plt.axis('off')
-        plt.show()
+        return self.generate_image()
 
 # Gradio Interface
 def train_snn(image_path, num_agents, epochs):
     snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(128, 128, 3), target_image_path=image_path)
-    snn.train(epochs=epochs)
-    generated_image = snn.generate_image()
-    return generated_image
+    generated_image = snn.train(epochs=epochs)
+    return (generated_image * 255).astype(np.uint8)
 
 interface = gr.Interface(
     fn=train_snn,