How to save the generated image singularly not as a grid at the end of training in this WGAN-GP Code

    def save_plt(self, epoch: int, results: np.ndarray):
        # construct an image from generated images with spacing between them using numpy
        w, h, c = results[0].shape
        # construct grind with self.grid_size
        grid = np.zeros((self.grid_size[0] * w + (self.grid_size[0] - 1) * self.spacing,
                         self.grid_size[1] * h + (self.grid_size[1] - 1) * self.spacing, c), dtype=np.uint8)
        for i in range(self.grid_size[0]):
            for j in range(self.grid_size[1]):
                grid[i * (w + self.spacing):i * (w + self.spacing) + w,
                j * (h + self.spacing):j * (h + self.spacing) + h] = results[i * self.grid_size[1] + j]

        grid = cv2.cvtColor(grid, cv2.COLOR_RGB2BGR)

        # save the image
        cv2.imwrite(f'{self.results_path}/img_{epoch}.png', grid)

        # save image to memory resized to gif size
        self.results.append(cv2.resize(grid, self.gif_size, interpolation=cv2.INTER_AREA))

    def on_epoch_end(self, epoch: int, logs: dict = None):
        # Define your custom code here that should be executed at the end of each epoch
        predictions = self.model.generator(self.seed, training=False)
        predictions_uint8 = (predictions * 127.5 + 127.5).numpy().astype(np.uint8)
        self.save_plt(epoch, predictions_uint8)

        if self.save_model:
            # save keras model to disk
            models_path = os.path.join(self.output_path, "model")
            os.makedirs(models_path, exist_ok=True)
            self.model.discriminator.save(models_path + "/discriminator.h5")
            self.model.generator.save(models_path + "/generator.h5")

    def on_train_end(self, logs: dict = None):
        # save the results as a gif with imageio

        # Create a list of imageio image objects from the OpenCV images
        # image is in BGR format, convert to RGB format when loading
        imageio_images = [imageio.core.util.Image(image[..., ::-1]) for image in self.results]

        # Write the imageio images to a GIF file
        imageio.mimsave(self.results_path + "/output.gif", imageio_images, duration=self.duration)

        # Save each image with a unique filename
        print(f"Generated images saved to: {self.results_path}")

here's the full code https://github.com/faiqalmdn/wgan-gp/blob/main/train.py

i've been tryin some loop code but the results still as an grid and still showing the images after every epoch

Upvotes: 0

Views: 14

Answers (0)

Related Questions