Update app.py

app.py

@@ -4,6 +4,10 @@ import numpy as np
 from PIL import Image
 from tqdm import tqdm
 import matplotlib.pyplot as plt
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
 
 # Disable GPU usage by default
 os.environ['CUDA_VISIBLE_DEVICES'] = ''
@@ -32,10 +36,10 @@ class SwarmAgent:
         self.velocity = velocity
 
 class SwarmNeuralNetwork:
-    def __init__(self, num_agents, image_shape, target_image_path):
+    def __init__(self, num_agents, image_shape, target_image):
         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.target_image = target_image
 
     def random_position(self):
         return np.random.randn(*self.image_shape)
@@ -43,10 +47,6 @@ class SwarmNeuralNetwork:
     def random_velocity(self):
         return np.random.randn(*self.image_shape) * 0.01
 
-    def load_target_image(self, img_path):
-        img = Image.open(img_path).convert('RGB').resize((self.image_shape[1], self.image_shape[0]))
-        return np.array(img) / 127.5 - 1
-
     def update_agents(self, timestep):
         for agent in self.agents:
             # Convert agent's position and target image into HDC space
@@ -74,18 +74,42 @@ class SwarmNeuralNetwork:
 
         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)
-    generated_image = snn.train(epochs=epochs)
-    return (generated_image * 255).astype(np.uint8)
+def preprocess_image(image):
+    """Preprocess the input image."""
+    if image is None:
+        raise ValueError("No image provided")
+    
+    if isinstance(image, np.ndarray):
+        # If it's already a numpy array, just resize and normalize
+        image = Image.fromarray(image)
+    elif isinstance(image, str):
+        # If it's a file path, open the image
+        image = Image.open(image)
+    else:
+        raise ValueError("Unsupported image type")
+    
+    image = image.convert('RGB').resize((128, 128))
+    return np.array(image) / 127.5 - 1
+
+def train_snn(image, num_agents, epochs):
+    try:
+        logging.info(f"Received image type: {type(image)}")
+        preprocessed_image = preprocess_image(image)
+        logging.info(f"Preprocessed image shape: {preprocessed_image.shape}")
+        
+        snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(128, 128, 3), target_image=preprocessed_image)
+        generated_image = snn.train(epochs=epochs)
+        return (generated_image * 255).astype(np.uint8)
+    except Exception as e:
+        logging.error(f"Error in train_snn: {str(e)}")
+        return None
 
 interface = gr.Interface(
     fn=train_snn,
     inputs=[
-        gr.Image(type="filepath", label="Upload Target Image"),
-        gr.Slider(minimum=100, maximum=1000, value=500, label="Number of Agents"),
-        gr.Slider(minimum=5, maximum=20, value=10, label="Number of Epochs")
+        gr.Image(type="numpy", label="Upload Target Image"),
+        gr.Slider(minimum=100, maximum=1000, value=500, step=50, label="Number of Agents"),
+        gr.Slider(minimum=5, maximum=20, value=10, step=1, label="Number of Epochs")
     ],
     outputs=gr.Image(type="numpy", label="Generated Image"),
     title="HDC Swarm Neural Network Image Generation"