From 63725a3e95bb1c1a657474e9ee6a2d9e1cd0a9ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Armando=20Ba=C3=B1uelos?=
 <32311654+abanuelo@users.noreply.github.com>
Date: Wed, 5 Jun 2024 11:49:00 -0500
Subject: [PATCH] feat: adding output gif for newtonian simulator (#273)

* fix: Adding modifications to codecov.yml

* feat: adding produced gif

* Adding colab

* fix: readding commented line

* feat: making gif export optional

* test: adding test case to test gif creation

* fix: adding graphical label

* fix: adding extra_gif options

* fix: reformatting files

---------

Co-authored-by: Armando Banuelos <abanuelo@Armandos-MacBook-Pro.local>
Co-authored-by: Armando Banuelos <abanuelo@Armandos-MBP.lan>
---
 src/scenic/simulators/newtonian/simulator.py | 27 +++++++++++++++++---
 tests/simulators/newtonian/driving.scenic    |  2 ++
 tests/simulators/newtonian/test_newtonian.py | 17 ++++++++++++
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/src/scenic/simulators/newtonian/simulator.py b/src/scenic/simulators/newtonian/simulator.py
index da453b0ab..fd38aa427 100644
--- a/src/scenic/simulators/newtonian/simulator.py
+++ b/src/scenic/simulators/newtonian/simulator.py
@@ -7,6 +7,9 @@
 import pathlib
 import time
 
+from PIL import Image
+import numpy as np
+
 import scenic.core.errors as errors  # isort: skip
 
 if errors.verbosityLevel == 0:  # suppress pygame advertisement at zero verbosity
@@ -55,21 +58,29 @@ class NewtonianSimulator(DrivingSimulator):
         when not otherwise specified is still 0.1 seconds.
     """
 
-    def __init__(self, network=None, render=False):
+    def __init__(self, network=None, render=False, export_gif=False):
         super().__init__()
+        self.export_gif = export_gif
         self.render = render
         self.network = network
 
     def createSimulation(self, scene, **kwargs):
-        return NewtonianSimulation(scene, self.network, self.render, **kwargs)
+        simulation = NewtonianSimulation(
+            scene, self.network, self.render, self.export_gif, **kwargs
+        )
+        if self.export_gif and self.render:
+            simulation.generate_gif("simulation.gif")
+        return simulation
 
 
 class NewtonianSimulation(DrivingSimulation):
     """Implementation of `Simulation` for the Newtonian simulator."""
 
-    def __init__(self, scene, network, render, timestep, **kwargs):
+    def __init__(self, scene, network, render, export_gif, timestep, **kwargs):
+        self.export_gif = export_gif
         self.render = render
         self.network = network
+        self.frames = []
 
         if timestep is None:
             timestep = 0.1
@@ -213,8 +224,18 @@ def draw_objects(self):
                 pygame.draw.polygon(self.screen, color, corners)
 
         pygame.display.update()
+
+        if self.export_gif:
+            frame = pygame.surfarray.array3d(self.screen)
+            frame = np.transpose(frame, (1, 0, 2))
+            self.frames.append(frame)
+
         time.sleep(self.timestep)
 
+    def generate_gif(self, filename="simulation.gif"):
+        imgs = [Image.fromarray(frame) for frame in self.frames]
+        imgs[0].save(filename, save_all=True, append_images=imgs[1:], duration=50, loop=0)
+
     def getProperties(self, obj, properties):
         yaw, _, _ = obj.parentOrientation.globalToLocalAngles(obj.heading, 0, 0)
 
diff --git a/tests/simulators/newtonian/driving.scenic b/tests/simulators/newtonian/driving.scenic
index 5fb4bd695..9a7b84d51 100644
--- a/tests/simulators/newtonian/driving.scenic
+++ b/tests/simulators/newtonian/driving.scenic
@@ -21,3 +21,5 @@ third = new Car on visible ego.road, with behavior Potpourri
 require abs((apparent heading of third) - 180 deg) <= 30 deg
 
 new Object visible, with width 0.1, with length 0.1
+
+terminate after 2 steps
diff --git a/tests/simulators/newtonian/test_newtonian.py b/tests/simulators/newtonian/test_newtonian.py
index c8dafae1d..a1ac5e4ad 100644
--- a/tests/simulators/newtonian/test_newtonian.py
+++ b/tests/simulators/newtonian/test_newtonian.py
@@ -1,5 +1,10 @@
+import os
+from pathlib import Path
+
+from PIL import Image as IPImage
 import pytest
 
+from scenic.domains.driving.roads import Network
 from scenic.simulators.newtonian import NewtonianSimulator
 from tests.utils import pickle_test, sampleScene, tryPickling
 
@@ -33,6 +38,18 @@ def check():
     check()  # If we fail here, something is leaking.
 
 
+@pytest.mark.graphical
+def test_gif_creation(loadLocalScenario):
+    scenario = loadLocalScenario("driving.scenic", mode2D=True)
+    scene, _ = scenario.generate(maxIterations=1000)
+    path = Path("assets") / "maps" / "CARLA" / "Town01.xodr"
+    network = Network.fromFile(path)
+    simulator = NewtonianSimulator(render=True, network=network, export_gif=True)
+    simulation = simulator.simulate(scene, maxSteps=100)
+    gif_path = Path("") / "simulation.gif"
+    assert os.path.exists(gif_path)
+
+
 @pickle_test
 def test_pickle(loadLocalScenario):
     scenario = tryPickling(loadLocalScenario("basic.scenic"))