Skip to content

Commit

Permalink
feat: adding output gif for newtonian simulator (#273)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
3 people authored Jun 5, 2024
1 parent 836b0f8 commit 63725a3
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 3 deletions.
27 changes: 24 additions & 3 deletions src/scenic/simulators/newtonian/simulator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down
2 changes: 2 additions & 0 deletions tests/simulators/newtonian/driving.scenic
Original file line number Diff line number Diff line change
Expand Up @@ -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
17 changes: 17 additions & 0 deletions tests/simulators/newtonian/test_newtonian.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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"))
Expand Down

0 comments on commit 63725a3

Please sign in to comment.