Create app.py

app.py

@@ -0,0 +1,199 @@
+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
+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)
+        
+        # Apply sharpening filter
+        image_pil = Image.fromarray((self.generated_image * 255).astype(np.uint8))
+        enhancer = ImageEnhance.Sharpness(image_pil)
+        self.generated_image = np.array(enhancer.enhance(2.0)) / 255.0
+
+    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, brightness, contrast, color):
+    snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(64, 64, 3), target_image=image)
+    
+    # Apply user-specified adjustments to the target image
+    image = ImageEnhance.Brightness(image).enhance(brightness)
+    image = ImageEnhance.Contrast(image).enhance(contrast)
+    image = ImageEnhance.Color(image).enhance(color)
+    
+    snn.target_image = snn.load_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"),
+        gr.Slider(minimum=0.5, maximum=2.0, value=1.0, label="Brightness"),
+        gr.Slider(minimum=0.5, maximum=2.0, value=1.0, label="Contrast"),
+        gr.Slider(minimum=0.5, maximum=2.0, value=1.0, label="Color Balance")
+    ],
+    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. Adjust brightness, contrast, and color balance for personalization."
+)
+
+interface.launch()