{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Import library","metadata":{}},{"cell_type":"markdown","source":"Get the dataset from kaggle (gan-getting-started)","metadata":{}},{"cell_type":"code","source":"import os\nimport torch\nimport torchvision\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torchvision.datasets import ImageFolder\nfrom torchvision import datasets, transforms, models\nfrom torchvision.utils import make_grid,save_image\nimport cv2\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import DataLoader, ConcatDataset\n%matplotlib inline","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-02-18T18:42:43.933822Z","iopub.execute_input":"2023-02-18T18:42:43.934127Z","iopub.status.idle":"2023-02-18T18:42:45.733288Z","shell.execute_reply.started":"2023-02-18T18:42:43.934047Z","shell.execute_reply":"2023-02-18T18:42:45.732536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(device)","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:45.735036Z","iopub.execute_input":"2023-02-18T18:42:45.735310Z","iopub.status.idle":"2023-02-18T18:42:45.806090Z","shell.execute_reply.started":"2023-02-18T18:42:45.735275Z","shell.execute_reply":"2023-02-18T18:42:45.804877Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prepare data","metadata":{}},{"cell_type":"code","source":"import os\nfrom PIL import Image\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\n\nclass Monet_Dataset(Dataset):\n    def __init__(self,\n                 pathToDataset = \"/kaggle/input/gan-getting-started/monet_jpg/\",\n                 transform = None\n                 ) -> None:\n        super().__init__()\n\n        self.pathToDataset = pathToDataset\n\n        # Defining transformer\n        if transform is not None:\n            self.transform = transform\n        else:\n            self.transform = transforms.ToTensor()\n\n        self.imageDir = os.listdir(pathToDataset)\n\n    def __len__(self):\n        return len(self.imageDir)\n\n    def __getitem__(self, index):\n        imagePath = self.pathToDataset + self.imageDir[index]\n        image = Image.open(imagePath)\n        image = self.transform(image)\n\n        return image","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:45.808074Z","iopub.execute_input":"2023-02-18T18:42:45.808474Z","iopub.status.idle":"2023-02-18T18:42:45.816979Z","shell.execute_reply.started":"2023-02-18T18:42:45.808433Z","shell.execute_reply":"2023-02-18T18:42:45.815567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Hyperparameter\nbatch_size = 8\nimage_size = 64\nnormalization_stats = (0.4819, 0.4324, 0.3845), (0.2604, 0.2520, 0.2539)\ndata_dir = r'/kaggle/input/gan-getting-started/monet_jpg/'\nnormal_dataset = Monet_Dataset(data_dir, transform=transforms.Compose([\n    transforms.Resize(image_size),\n    transforms.ToTensor(),\n    transforms.Normalize(*normalization_stats)]))","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:45.820289Z","iopub.execute_input":"2023-02-18T18:42:45.820651Z","iopub.status.idle":"2023-02-18T18:42:45.873226Z","shell.execute_reply.started":"2023-02-18T18:42:45.820612Z","shell.execute_reply":"2023-02-18T18:42:45.872456Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_loader = DataLoader(dataset=normal_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=False)","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:45.876036Z","iopub.execute_input":"2023-02-18T18:42:45.876326Z","iopub.status.idle":"2023-02-18T18:42:45.883033Z","shell.execute_reply.started":"2023-02-18T18:42:45.876293Z","shell.execute_reply":"2023-02-18T18:42:45.882179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Display some image","metadata":{}},{"cell_type":"code","source":"def denorm(image):\n    return image * normalization_stats[1][0] + normalization_stats[0][0]","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:45.884381Z","iopub.execute_input":"2023-02-18T18:42:45.885230Z","iopub.status.idle":"2023-02-18T18:42:45.889583Z","shell.execute_reply.started":"2023-02-18T18:42:45.885195Z","shell.execute_reply":"2023-02-18T18:42:45.888653Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_images(images, nmax=64):\n    fig, ax = plt.subplots(figsize=(20, 20))\n    ax.set_xticks([]); ax.set_yticks([])\n    ax.imshow(make_grid(denorm(images.detach()[:nmax]), nrow=8).permute(1, 2, 0))\n    \ndef show_batch(dataloader, nmax=64):\n    for images in dataloader:\n        show_images(images, nmax)\n        break","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:45.890872Z","iopub.execute_input":"2023-02-18T18:42:45.891269Z","iopub.status.idle":"2023-02-18T18:42:45.898485Z","shell.execute_reply.started":"2023-02-18T18:42:45.891235Z","shell.execute_reply":"2023-02-18T18:42:45.897778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_batch(data_loader)","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:45.899792Z","iopub.execute_input":"2023-02-18T18:42:45.900244Z","iopub.status.idle":"2023-02-18T18:42:46.439836Z","shell.execute_reply.started":"2023-02-18T18:42:45.900210Z","shell.execute_reply":"2023-02-18T18:42:46.439149Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"discriminator = nn.Sequential(\n    # Input is 3 x 64 x 64\n    nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1, bias=False),\n    nn.BatchNorm2d(64),\n    nn.LeakyReLU(0.2, inplace=True),\n    # Layer Output: 64 x 32 x 32\n    \n    nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),\n    nn.BatchNorm2d(128),\n    nn.LeakyReLU(0.2, inplace=True),\n    # Layer Output: 128 x 16 x 16\n    \n    nn.Conv2d(128, 128, kernel_size=4, stride=2, padding=1, bias=False),\n    nn.BatchNorm2d(128),\n    nn.LeakyReLU(0.2, inplace=True),\n    # Layer Output: 128 x 8 x 8\n    \n    nn.Conv2d(128, 128, kernel_size=4, stride=2, padding=1, bias=False),\n    nn.BatchNorm2d(128),\n    nn.LeakyReLU(0.2, inplace=True),\n    # Layer Output: 128 x 4 x 4\n    \n    # With a 4x4, we can condense the channels into a 1 x 1 x 1 to produce output\n    nn.Conv2d(128, 1, kernel_size=4, stride=1, padding=0, bias=False),\n    nn.Flatten(),\n    nn.Sigmoid()\n)\ndiscriminator.to(device)","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:46.441096Z","iopub.execute_input":"2023-02-18T18:42:46.441356Z","iopub.status.idle":"2023-02-18T18:42:49.258796Z","shell.execute_reply.started":"2023-02-18T18:42:46.441326Z","shell.execute_reply":"2023-02-18T18:42:49.258022Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed_size = 16","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:49.261705Z","iopub.execute_input":"2023-02-18T18:42:49.261936Z","iopub.status.idle":"2023-02-18T18:42:49.265718Z","shell.execute_reply.started":"2023-02-18T18:42:49.261904Z","shell.execute_reply":"2023-02-18T18:42:49.264763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generator = nn.Sequential(\n    # Input seed_size x 1 x 1\n    nn.ConvTranspose2d(seed_size, 128, kernel_size=4, padding=0, stride=1, bias=False),\n    nn.BatchNorm2d(128),\n    nn.ReLU(True),\n    # Layer output: 256 x 4 x 4\n    \n    nn.ConvTranspose2d(128, 128, kernel_size=4, padding=1, stride=2, bias=False),\n    nn.BatchNorm2d(128),\n    nn.ReLU(True),\n    # Layer output: 128 x 8 x 8\n    \n    nn.ConvTranspose2d(128, 128, kernel_size=4, padding=1, stride=2, bias=False),\n    nn.BatchNorm2d(128),\n    nn.ReLU(True),\n    # Layer output: 64 x 16 x 16\n    \n    nn.ConvTranspose2d(128, 64, kernel_size=4, padding=1, stride=2, bias=False),\n    nn.BatchNorm2d(64),\n    nn.ReLU(True),\n    # Layer output: 32 x 32 x 32\n    \n    nn.ConvTranspose2d(64, 3, kernel_size=4, padding=1, stride=2, bias=False),\n    nn.Tanh()\n    # Output: 3 x 64 x 64\n)\ngenerator.to(device)","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:49.267032Z","iopub.execute_input":"2023-02-18T18:42:49.267519Z","iopub.status.idle":"2023-02-18T18:42:49.289333Z","shell.execute_reply.started":"2023-02-18T18:42:49.267485Z","shell.execute_reply":"2023-02-18T18:42:49.288738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"def train_discriminator(real_pokemon, disc_optimizer):\n    # Reset the gradients for the optimizer\n    disc_optimizer.zero_grad()\n    \n    # Train on the real images\n    real_predictions = discriminator(real_pokemon)\n    # real_targets = torch.zeros(real_pokemon.size(0), 1, device=device) # All of these are real, so the target is 0.\n    real_targets = torch.rand(real_pokemon.size(0), 1, device=device) * (0.1 - 0) + 0 # Add some noisy labels to make the discriminator think harder.\n    real_loss = F.binary_cross_entropy(real_predictions, real_targets) # Can do binary loss function because it is a binary classifier\n    real_score = torch.mean(real_predictions).item() # How well does the discriminator classify the real pokemon? (Higher score is better for the discriminator)\n    \n    # Make some latent tensors to seed the generator\n    latent_batch = torch.randn(batch_size, seed_size, 1, 1, device=device)\n    \n    # Get some fake pokemon\n    fake_pokemon = generator(latent_batch)\n    \n    # Train on the generator's current efforts to trick the discriminator\n    gen_predictions = discriminator(fake_pokemon)\n    # gen_targets = torch.ones(fake_pokemon.size(0), 1, device=device)\n    gen_targets = torch.rand(fake_pokemon.size(0), 1, device=device) * (1 - 0.9) + 0.9 # Add some noisy labels to make the discriminator think harder.\n    gen_loss = F.binary_cross_entropy(gen_predictions, gen_targets)\n    gen_score = torch.mean(gen_predictions).item() # How well did the discriminator classify the fake pokemon? (Lower score is better for the discriminator)\n    \n    # Update the discriminator weights\n    total_loss = real_loss + gen_loss\n    total_loss.backward()\n    disc_optimizer.step()\n    return total_loss.item(), real_score, gen_score","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:49.291062Z","iopub.execute_input":"2023-02-18T18:42:49.291512Z","iopub.status.idle":"2023-02-18T18:42:49.298944Z","shell.execute_reply.started":"2023-02-18T18:42:49.291477Z","shell.execute_reply":"2023-02-18T18:42:49.298282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_generator(gen_optimizer):\n    # Clear the generator gradients\n    gen_optimizer.zero_grad()\n    \n    # Generate some fake pokemon\n    latent_batch = torch.randn(batch_size, seed_size, 1, 1, device=device)\n    fake_pokemon = generator(latent_batch)\n    \n    # Test against the discriminator\n    disc_predictions = discriminator(fake_pokemon)\n    targets = torch.zeros(fake_pokemon.size(0), 1, device=device) # We want the discriminator to think these images are real.\n    loss = F.binary_cross_entropy(disc_predictions, targets) # How well did the generator do? (How much did the discriminator believe the generator?)\n    \n    # Update the generator based on how well it fooled the discriminator\n    loss.backward()\n    gen_optimizer.step()\n    \n    # Return generator loss\n    return loss.item()","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:49.300155Z","iopub.execute_input":"2023-02-18T18:42:49.300816Z","iopub.status.idle":"2023-02-18T18:42:49.310615Z","shell.execute_reply.started":"2023-02-18T18:42:49.300779Z","shell.execute_reply":"2023-02-18T18:42:49.309885Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Save result","metadata":{}},{"cell_type":"code","source":"import os\nfrom torchvision.utils import save_image\n\nRESULTS_DIR = 'results'\nos.makedirs(RESULTS_DIR, exist_ok=True)\n\ndef save_results(index, latent_batch, show=True):\n    # Generate fake pokemon\n    fake_pokemon = generator(latent_batch)\n    \n    # Make the filename for the output\n    fake_file = \"result-image-{0:0=4d}.png\".format(index)\n    \n    # Save the image\n    save_image(denorm(fake_pokemon), os.path.join(RESULTS_DIR, fake_file), nrow=8)\n    print(\"Result Saved!\")\n    \n    if show:\n        fig, ax = plt.subplots(figsize=(20, 20))\n        ax.set_xticks([]); ax.set_yticks([])\n        ax.imshow(make_grid(fake_pokemon.cpu().detach(), nrow=8).permute(1, 2, 0))","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:49.311988Z","iopub.execute_input":"2023-02-18T18:42:49.312274Z","iopub.status.idle":"2023-02-18T18:42:49.321323Z","shell.execute_reply.started":"2023-02-18T18:42:49.312240Z","shell.execute_reply":"2023-02-18T18:42:49.320703Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm.notebook import tqdm\nimport torch.nn.functional as F\n\n# Static generation seed batch\nfixed_latent_batch = torch.randn(64, seed_size, 1, 1, device=device)\n\ndef train(epochs, learning_rate, start_idx=1):\n    # Empty the GPU cache to save some memory\n    torch.cuda.empty_cache()\n    \n    # Track losses and scores\n    disc_losses = []\n    disc_scores = []\n    gen_losses = []\n    gen_scores = []\n    \n    # Create the optimizers\n    disc_optimizer = torch.optim.Adam(discriminator.parameters(), lr=learning_rate, betas=(0.5, 0.9))\n    gen_optimizer = torch.optim.Adam(generator.parameters(), lr=learning_rate, betas=(0.5, 0.9))\n    \n    # Run the loop\n    for epoch in range(epochs):\n        # Go through each image\n        for real_img in tqdm(data_loader):\n            real_img = real_img.to(device)\n            # Train the discriminator\n            disc_loss, real_score, gen_score = train_discriminator(real_img, disc_optimizer)\n\n            # Train the generator\n            gen_loss = train_generator(gen_optimizer)\n        \n        # Collect results\n        disc_losses.append(disc_loss)\n        disc_scores.append(real_score)\n        gen_losses.append(gen_loss)\n        gen_scores.append(gen_score)\n        \n        # Print the losses and scores\n        print(\"Epoch [{}/{}], gen_loss: {:.4f}, disc_loss: {:.4f}, real_score: {:.4f}, gen_score: {:.4f}\".format(\n            epoch+start_idx, epochs, gen_loss, disc_loss, real_score, gen_score))\n        \n        # Save the images and show the progress\n        save_results(epoch + start_idx, fixed_latent_batch, show=False)\n    \n    # Return stats\n    return disc_losses, disc_scores, gen_losses, gen_scores","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:49.323466Z","iopub.execute_input":"2023-02-18T18:42:49.324286Z","iopub.status.idle":"2023-02-18T18:42:49.336929Z","shell.execute_reply.started":"2023-02-18T18:42:49.323633Z","shell.execute_reply":"2023-02-18T18:42:49.336124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"learning_rate = 1e-4\nepochs = 300","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:42:49.338317Z","iopub.execute_input":"2023-02-18T18:42:49.338563Z","iopub.status.idle":"2023-02-18T18:42:49.343948Z","shell.execute_reply.started":"2023-02-18T18:42:49.338531Z","shell.execute_reply":"2023-02-18T18:42:49.343196Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"history = train(epochs, learning_rate)","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-02-18T18:42:49.345475Z","iopub.execute_input":"2023-02-18T18:42:49.345787Z","iopub.status.idle":"2023-02-18T18:48:56.838839Z","shell.execute_reply.started":"2023-02-18T18:42:49.345738Z","shell.execute_reply":"2023-02-18T18:48:56.837856Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Show some fake images","metadata":{}},{"cell_type":"code","source":"from IPython.display import Image\nImage('./results/result-image-0010.png')","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:49:18.113943Z","iopub.execute_input":"2023-02-18T18:49:18.114228Z","iopub.status.idle":"2023-02-18T18:49:18.138771Z","shell.execute_reply.started":"2023-02-18T18:49:18.114194Z","shell.execute_reply":"2023-02-18T18:49:18.138000Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Image('./results/result-image-0025.png')","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:49:22.900618Z","iopub.execute_input":"2023-02-18T18:49:22.901306Z","iopub.status.idle":"2023-02-18T18:49:22.919260Z","shell.execute_reply.started":"2023-02-18T18:49:22.901273Z","shell.execute_reply":"2023-02-18T18:49:22.918516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Image('./results/result-image-0050.png')","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:49:23.596563Z","iopub.execute_input":"2023-02-18T18:49:23.596829Z","iopub.status.idle":"2023-02-18T18:49:23.613487Z","shell.execute_reply.started":"2023-02-18T18:49:23.596798Z","shell.execute_reply":"2023-02-18T18:49:23.612721Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Image('./results/result-image-0150.png')","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:49:24.165972Z","iopub.execute_input":"2023-02-18T18:49:24.166257Z","iopub.status.idle":"2023-02-18T18:49:24.183585Z","shell.execute_reply.started":"2023-02-18T18:49:24.166223Z","shell.execute_reply":"2023-02-18T18:49:24.183022Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Image('./results/result-image-0300.png')","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:49:24.635648Z","iopub.execute_input":"2023-02-18T18:49:24.636208Z","iopub.status.idle":"2023-02-18T18:49:24.651415Z","shell.execute_reply.started":"2023-02-18T18:49:24.636145Z","shell.execute_reply":"2023-02-18T18:49:24.650637Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"disc_losses, disc_scores, gen_losses, gen_scores = history","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:48:56.948446Z","iopub.execute_input":"2023-02-18T18:48:56.949123Z","iopub.status.idle":"2023-02-18T18:48:56.952884Z","shell.execute_reply.started":"2023-02-18T18:48:56.949092Z","shell.execute_reply":"2023-02-18T18:48:56.952105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Plotting","metadata":{}},{"cell_type":"code","source":"# Plot generator and discriminator losses\nplt.plot(disc_losses, '-')\nplt.plot(gen_losses, '-')\nplt.xlabel('Epoch')\nplt.ylabel('Loss')\nplt.legend(['Discriminator', 'Generator'])\nplt.title('Losses');","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:49:47.575003Z","iopub.execute_input":"2023-02-18T18:49:47.575762Z","iopub.status.idle":"2023-02-18T18:49:47.844437Z","shell.execute_reply.started":"2023-02-18T18:49:47.575703Z","shell.execute_reply":"2023-02-18T18:49:47.840972Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plots scores vs. epochs\nplt.plot(disc_scores, '-')\nplt.plot(gen_scores, '-')\nplt.xlabel('Epoch')\nplt.ylabel('Score')\nplt.legend(['Real', 'Fake'])\nplt.title('Scores');","metadata":{"execution":{"iopub.status.busy":"2023-02-18T18:49:48.291753Z","iopub.execute_input":"2023-02-18T18:49:48.292021Z","iopub.status.idle":"2023-02-18T18:49:48.540751Z","shell.execute_reply.started":"2023-02-18T18:49:48.291991Z","shell.execute_reply":"2023-02-18T18:49:48.540045Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%matplotlib notebook\n\nimport random\nimport os\n\nimport numpy as np\nfrom PIL import Image\n\nimport matplotlib\nimport matplotlib.pyplot as plt\n\nimport matplotlib.animation as animation\n\n\nfps = 15\nnSeconds = 10\n\npath = \"/kaggle/working/results/\"\nimages = os.listdir(path)\nimages.sort()\n\n# print(images)\n\nsnapshots = [Image.open(path + i) for i in images]\nprint(len(snapshots))\nfor _ in range(500):\n    if _ % 2 == 0:\n        snapshots.append(Image.open(path + images[-1]))\n    else:\n        snapshots.append(Image.open(path + images[-2]))\nprint(len(snapshots))\n\n# First set up the figure, the axis, and the plot element we want to animate\nfig = plt.figure( figsize=(8,8) )\nim = plt.imshow(a, interpolation='none', aspect='auto', vmin=0, vmax=1)\n\ndef animate_func(i):\n    if i % fps == 0:\n        print( '.', end ='' )\n\n    im.set_array(snapshots[i])\n    return [im]\n\nanim = animation.FuncAnimation(\n                               fig, \n                               animate_func, \n                               frames = nSeconds * fps,\n                               interval = 1000 / fps, # in ms\n                               )\n\nanim.save('test_anim.mp4', fps=fps) # , extra_args=['-vcodec', 'libx264'])\n\nprint('Done!')","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-02-20T07:15:03.802626Z","iopub.execute_input":"2023-02-20T07:15:03.802968Z","iopub.status.idle":"2023-02-20T07:15:15.374383Z","shell.execute_reply.started":"2023-02-20T07:15:03.802923Z","shell.execute_reply":"2023-02-20T07:15:15.373122Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}