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 from tqdm import tqdm class SwarmAgent: def __init__(self, position, velocity): self.position = position self.velocity = velocity self.m = np.zeros_like(position) self.v = np.zeros_like(position) class SwarmNeuralNetwork: def __init__(self, num_agents, image_shape, target_image): self.image_shape = image_shape self.resized_shape = (64, 64, 3) self.agents = [SwarmAgent(self.random_position(), self.random_velocity()) for _ in range(num_agents)] self.target_image = self.load_target_image(target_image) self.generated_image = np.random.randn(*image_shape) # Start with noise self.mobilenet = self.load_mobilenet_model() self.current_epoch = 0 self.noise_schedule = np.linspace(0.1, 0.002, 1000) # Noise schedule def random_position(self): return np.random.randn(*self.image_shape) # Use Gaussian noise def random_velocity(self): return np.random.randn(*self.image_shape) * 0.01 def load_target_image(self, img): img = img.resize((self.image_shape[1], self.image_shape[0])) img_array = np.array(img) / 127.5 - 1 # Normalize to [-1, 1] plt.imshow((img_array + 1) / 2) # Convert back to [0, 1] for display plt.title('Target Image') plt.show() return img_array def resize_image(self, image): return resize(image, self.resized_shape, anti_aliasing=True) def load_mobilenet_model(self): mobilenet = MobileNetV2(weights='imagenet', include_top=False, input_shape=self.resized_shape) return Model(inputs=mobilenet.input, outputs=mobilenet.get_layer('block_13_expand_relu').output) def add_positional_encoding(self, image): h, w, c = image.shape pos_enc = np.zeros_like(image) for i in range(h): for j in range(w): pos_enc[i, j, :] = [i/h, j/w, 0] return image + pos_enc def multi_head_attention(self, agent, num_heads=4): attention_scores = [] for _ in range(num_heads): similarity = np.exp(-np.sum((agent.position - self.target_image)**2, axis=-1)) attention_score = similarity / np.sum(similarity) attention_scores.append(attention_score) attention = np.mean(attention_scores, axis=0) return np.expand_dims(attention, axis=-1) def multi_scale_perceptual_loss(self, agent_positions): target_image_resized = self.resize_image((self.target_image + 1) / 2) # Convert to [0, 1] for MobileNet target_image_preprocessed = preprocess_input(target_image_resized[np.newaxis, ...] * 255) # MobileNet expects [0, 255] target_features = self.mobilenet.predict(target_image_preprocessed) losses = [] for agent_position in agent_positions: agent_image_resized = self.resize_image((agent_position + 1) / 2) agent_image_preprocessed = preprocess_input(agent_image_resized[np.newaxis, ...] * 255) agent_features = self.mobilenet.predict(agent_image_preprocessed) loss = np.mean((target_features - agent_features)**2) losses.append(1 / (1 + loss)) return np.array(losses) def update_agents(self, timestep): noise_level = self.noise_schedule[min(timestep, len(self.noise_schedule) - 1)] for agent in self.agents: # Predict noise predicted_noise = agent.position - self.target_image # Denoise denoised = (agent.position - noise_level * predicted_noise) / (1 - noise_level) # Add scaled noise for next step agent.position = denoised + np.random.randn(*self.image_shape) * np.sqrt(noise_level) # Clip values agent.position = np.clip(agent.position, -1, 1) def generate_image(self): self.generated_image = np.mean([agent.position for agent in self.agents], axis=0) # Normalize to [0, 1] range for display self.generated_image = (self.generated_image + 1) / 2 self.generated_image = np.clip(self.generated_image, 0, 1) def train(self, epochs): logging.basicConfig(filename='training.log', level=logging.INFO) for epoch in tqdm(range(epochs), desc="Training Epochs"): self.update_agents(epoch) self.generate_image() mse = np.mean(((self.generated_image * 2 - 1) - self.target_image)**2) logging.info(f"Epoch {epoch}, MSE: {mse}") if epoch % 10 == 0: print(f"Epoch {epoch}, MSE: {mse}") self.display_image(self.generated_image, title=f'Epoch {epoch}') self.current_epoch += 1 def display_image(self, image, title=''): plt.imshow(image) plt.title(title) plt.axis('off') plt.show() def display_agent_positions(self, epoch): fig, ax = plt.subplots() positions = np.array([agent.position for agent in self.agents]) ax.imshow(self.generated_image, extent=[0, self.image_shape[1], 0, self.image_shape[0]]) ax.scatter(positions[:, :, 0].flatten(), positions[:, :, 1].flatten(), s=1, c='red') plt.title(f'Agent Positions at Epoch {epoch}') plt.show() def save_model(self, filename): model_state = { 'agents': self.agents, 'generated_image': self.generated_image, 'current_epoch': self.current_epoch } np.save(filename, model_state) def load_model(self, filename): model_state = np.load(filename, allow_pickle=True).item() self.agents = model_state['agents'] self.generated_image = model_state['generated_image'] self.current_epoch = model_state['current_epoch'] def generate_new_image(self, num_steps=1000): for agent in self.agents: agent.position = np.random.randn(*self.image_shape) for step in tqdm(range(num_steps), desc="Generating Image"): self.update_agents(num_steps - step - 1) # Reverse order self.generate_image() return self.generated_image # Gradio Interface def train_snn(image, num_agents, epochs): snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(64, 64, 3), target_image=image) snn.train(epochs=epochs) snn.save_model('snn_model.npy') return snn.generated_image def generate_new_image(): snn = SwarmNeuralNetwork(num_agents=2000, image_shape=(64, 64, 3), target_image=None) snn.load_model('snn_model.npy') new_image = snn.generate_new_image() return new_image interface = gr.Interface( fn=train_snn, inputs=[ gr.Image(type="pil", label="Upload Target Image"), gr.Slider(minimum=500, maximum=3000, value=2000, label="Number of Agents"), gr.Slider(minimum=10, maximum=200, value=100, label="Number of Epochs") ], outputs=gr.Image(type="numpy", label="Generated Image"), title="Swarm Neural Network Image Generation", description="Upload an image and set the number of agents and epochs to train the Swarm Neural Network to generate a new image." ) interface.launch()