Update app.py

app.py

@@ -6,7 +6,7 @@ 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, ImageEnhance
 from tqdm import tqdm
 
 class SwarmAgent:
@@ -19,7 +19,7 @@ class SwarmAgent:
 class SwarmNeuralNetwork:
     def __init__(self, num_agents, image_shape, target_image):
         self.image_shape = image_shape
-        self.resized_shape = (128, 128, 3)  # Increased resolution
+        self.resized_shape = (256, 256, 3)  # Increased resolution
         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
@@ -102,12 +102,6 @@ class SwarmNeuralNetwork:
         # 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))
-        image_pil = image_pil.filter(ImageFilter.SHARPEN)
-        image_pil = image_pil.filter(ImageFilter.DETAIL)
-        self.generated_image = np.array(image_pil) / 255.0
 
     def train(self, epochs):
         logging.basicConfig(filename='training.log', level=logging.INFO)
@@ -162,9 +156,19 @@ class SwarmNeuralNetwork:
         self.generate_image()
         return self.generated_image
 
+    def apply_super_resolution(self, image):
+        import cv2
+        sr = cv2.dnn_superres.DnnSuperResImpl_create()
+        path = "EDSR_x3.pb"  # Path to a pre-trained super-resolution model file (download it from OpenCV's repository)
+        sr.readModel(path)
+        sr.setModel("edsr", 3)  # Use EDSR model with a scale factor of 3
+        image = (image * 255).astype(np.uint8)
+        upscaled = sr.upsample(image)
+        return upscaled / 255.0
+
 # Gradio Interface
 def train_snn(image, num_agents, epochs, brightness, contrast, color):
-    snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(128, 128, 3), target_image=image)
+    snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(256, 256, 3), target_image=image)
     
     # Apply user-specified adjustments to the target image
     image = ImageEnhance.Brightness(image).enhance(brightness)
@@ -174,13 +178,16 @@ def train_snn(image, num_agents, epochs, brightness, contrast, color):
     snn.target_image = snn.load_target_image(image)
     snn.train(epochs=epochs)
     snn.save_model('snn_model.npy')
-    return snn.generated_image
+    generated_image = snn.generated_image
+    upscaled_image = snn.apply_super_resolution(generated_image)
+    return upscaled_image
 
 def generate_new_image():
-    snn = SwarmNeuralNetwork(num_agents=2000, image_shape=(128, 128, 3), target_image=None)
+    snn = SwarmNeuralNetwork(num_agents=2000, image_shape=(256, 256, 3), target_image=None)
     snn.load_model('snn_model.npy')
     new_image = snn.generate_new_image()
-    return new_image
+    upscaled_image = snn.apply_super_resolution(new_image)
+    return upscaled_image
 
 interface = gr.Interface(
     fn=train_snn,
@@ -197,4 +204,4 @@ interface = gr.Interface(
     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()
+interface.launch()