Fix TextureAtlas issues (#1949)

* Textures are now stored using the `Texture.atlas_name` meaning we only store textures with different hash-orientation combo. It doesn't care what hitbox algo the texture has. This mean the number of textures in the atlas are always the same as the uv data simplifying things.
* Resizing the atlas should no longer clear the image counter
* Rebuilding the atlas should now clear the image counter
* The atlas will now smartly rebuild itself if images have been removed before attempting to resize
* Added an initial ABC class for atlases
* Added sanity check in image ref counter raising error on underflow
* Fixed the experimenal atlas load/save code
* Added various comments to make it easier to inspect the code later

arcade/context.py

@@ -43,7 +43,7 @@ class ArcadeContext(Context):
                         it's not clear what thread will gc the object.
     """
 
-    atlas_size = 512, 512
+    atlas_size: Tuple[int, int] = 512, 512
 
     def __init__(self, window: pyglet.window.Window, gc_mode: str = "context_gc", gl_api: str = "gl"):
 

arcade/experimental/atlas_load_save.py

@@ -1,5 +1,7 @@
 """
 Quick and dirty atlas load/save testing.
+Loading and saving atlases are not officially supported.
+This is simply an experiment.
 
 Dump atlas:
 python arcade/experimental/atlas_load_save.py save
@@ -19,7 +21,7 @@
 import arcade
 from arcade.texture_atlas.helpers import save_atlas, load_atlas
 
-MODE = 'load'
+MODE = 'save'
 RESOURCE_ROOT = arcade.resources.ASSET_PATH
 DESTINATION = Path.cwd()
 
@@ -78,7 +80,7 @@ def __init__(self):
 
         # Make a sprite for each texture
         self.sp = arcade.SpriteList(atlas=self.atlas)
-        for i, texture in enumerate(self.atlas._textures):
+        for i, texture in enumerate(self.atlas.textures):
             pos = i * 64
             sprite = arcade.Sprite(
                 texture,

arcade/texture_atlas/base.py

@@ -10,9 +10,21 @@
 Pyglet atlases are located here:
 https://github.com/einarf/pyglet/blob/master/pyglet/image/atlas.py
 
+Allocation:
+Pyglet's allocator is a simple row based allocator only keeping
+track of horizontal strips and how far in the x direction the
+each strip is filled. We can't really "deallocate" unless it's
+a region at the end of a strip and even doing that is awkward.
+
+When an image is removed from the atlas we simply just lose that
+region until we rebuild the atlas. It can be a good idea to count
+the number of lost regions to use as an indicator later. When an
+atlas is full we can first rebuild it if there are lost regions
+instead of increasing the size.
 """
 from __future__ import annotations
 
+import abc
 import copy
 import math
 import time
@@ -30,7 +42,7 @@
 from array import array
 from collections import deque
 from contextlib import contextmanager
-from weakref import WeakSet
+from weakref import WeakSet, WeakValueDictionary
 
 import PIL
 import PIL.Image
@@ -161,33 +173,90 @@ class ImageDataRefCounter:
     """
     Helper class to keep track of how many times an image is used
     by a texture in the atlas to determine when it's safe to remove it.
+
+    Multiple Texture instances can and will use the same ImageData
+    instance.
     """
     def __init__(self) -> None:
         self._data: Dict[str, int] = {}
+        self._num_decref = 0
 
     def inc_ref(self, image_data: "ImageData") -> None:
+        """Increment the reference count for an image."""
         self._data[image_data.hash] = self._data.get(image_data.hash, 0) + 1
 
-    def dec_ref(self, image_data: "ImageData") -> None:
-        # TODO: Should we raise an error if the ref count is 0?
+    def dec_ref(self, image_data: "ImageData") -> int:
+        """
+        Decrement the reference count for an image returning the new value.
+        """
         if image_data.hash not in self._data:
-            return
-        self._data[image_data.hash] -= 1
-        if self._data[image_data.hash] == 0:
+            raise RuntimeError(f"Image {image_data.hash} not in ref counter")
+
+        val = self._data[image_data.hash] - 1
+        self._data[image_data.hash] = val
+
+        if val < 0:
+            raise RuntimeError(f"Image {image_data.hash} ref count went below zero")
+        if val == 0:
             del self._data[image_data.hash]
 
-    def get_refs(self, image_data: "ImageData") -> int:
+        self._num_decref += 1
+
+        return val
+
+    def get_ref_count(self, image_data: "ImageData") -> int:
+        """
+        Get the reference count for an image.
+
+        Args:
+            image_data (ImageData): The image to get the reference count for
+        """
         return self._data.get(image_data.hash, 0)
 
-    def count_refs(self) -> int:
+    def count_all_refs(self) -> int:
         """Helper function to count the total number of references."""
         return sum(self._data.values())
 
+    def get_total_decref(self, reset=True) -> int:
+        """
+        Get the total number of decrefs.
+
+        Args:
+            reset (bool): Reset the counter after getting the value
+        """
+        num_decref = self._num_decref
+        if reset:
+            self._num_decref = 0
+        return num_decref
+
     def __len__(self) -> int:
         return len(self._data)
 
+    def __repr__(self) -> str:
+        return f"<ImageDataRefCounter ref_count={self.count_all_refs()} data={self._data}>"
+
+
+class BaseTextureAtlas(abc.ABC):
 
-class TextureAtlas:
+    def __init__(self, ctx: Optional["ArcadeContext"]):
+        self._ctx = ctx or arcade.get_window().ctx
+
+    @property
+    def ctx(self) -> "ArcadeContext":
+        return self._ctx
+
+    @abc.abstractmethod
+    def add(self, texture: "Texture") -> Tuple[int, AtlasRegion]:
+        """Add a texture to the atlas."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def remove(self, texture: "Texture") -> None:
+        """Remove a texture from the atlas."""
+        raise NotImplementedError
+
+
+class TextureAtlas(BaseTextureAtlas):
     """
     A texture atlas with a size in a context.
 
@@ -230,7 +299,7 @@ def __init__(
         ctx: Optional["ArcadeContext"] = None,
         capacity: int = 2,
     ):
-        self._ctx = ctx or arcade.get_window().ctx
+        super().__init__(ctx)
         self._max_size = self._ctx.info.MAX_VIEWPORT_DIMS
         self._size: Tuple[int, int] = size
         self._allocator = Allocator(*self._size)
@@ -270,10 +339,11 @@ def __init__(
         # when to remove an image from the atlas.
         self._image_ref_count = ImageDataRefCounter()
 
-        # A list of all the images this atlas contains
+        # A list of all the images this atlas contains.
+        # Unique by: Internal hash property
         self._images: WeakSet[ImageData] = WeakSet()
-        # A set of textures this atlas contains for fast lookups + set operations
-        self._textures: WeakSet["Texture"] = WeakSet()
+        # atlas_name: Texture
+        self._textures: WeakValueDictionary[str, "Texture"] = WeakValueDictionary()
 
         # Texture containing texture coordinates for images and textures
         # The 4096 width is a safe constant for all GL implementations
@@ -414,7 +484,7 @@ def textures(self) -> List["Texture"]:
         A new list is constructed from the internal weak set of textures.
 
         """
-        return list(self._textures)
+        return list(self._textures.values())
 
     @property
     def images(self) -> List["ImageData"]:
@@ -432,6 +502,7 @@ def add(self, texture: "Texture") -> Tuple[int, AtlasRegion]:
 
         :param texture: The texture to add
         :return: texture_id, AtlasRegion tuple
+        :raises AllocatorException: If there are no room for the texture
         """
         # If the texture is already in the atlas we also have the image
         # and can return early with the texture id and region
@@ -442,28 +513,36 @@ def add(self, texture: "Texture") -> Tuple[int, AtlasRegion]:
 
         LOG.info("Attempting to add texture: %s", texture.atlas_name)
 
-        # Add the image if we don't already have it.
-        # If the atlas is full we will try to resize it.
+        # Add the *image* to the atlas if it's not already there
         if not self.has_image(texture.image_data):
             try:
                 x, y, slot, region = self.allocate(texture.image_data)
             except AllocatorException:
                 LOG.info("[%s] No room for %s size %s", id(self), texture.atlas_name, texture.image.size)
-                if self._auto_resize:
-                    width = min(self.width * 2, self.max_width)
-                    height = min(self.height * 2, self.max_height)
-                    if self._size == (width, height):
-                        raise
-                    self.resize((width, height))
+                if not self._auto_resize:
+                    raise
+
+                # If we have lost regions we can try to rebuild the atlas
+                removed_image_count = self._image_ref_count.get_total_decref()
+                if removed_image_count > 0:
+                    LOG.info("[%s] Rebuilding atlas due to %s lost images", id(self), removed_image_count)
+                    self.rebuild()
                     return self.add(texture)
-                else:
+
+                width = min(self.width * 2, self.max_width)
+                height = min(self.height * 2, self.max_height)
+                if self._size == (width, height):
                     raise
+                self.resize((width, height))
+                return self.add(texture)
 
+            # Write the pixel data to the atlas texture
             self.write_image(texture.image_data.image, x, y)
 
+        info = self._allocate_texture(texture)
         self._image_ref_count.inc_ref(texture.image_data)
         texture.add_atlas_ref(self)
-        return self._allocate_texture(texture)
+        return info
 
     def _allocate_texture(self, texture: "Texture") -> Tuple[int, AtlasRegion]:
         """
@@ -480,14 +559,14 @@ def _allocate_texture(self, texture: "Texture") -> Tuple[int, AtlasRegion]:
                 f"Max number of slots: {self._num_texture_slots}"
             ))
 
-        # Add the texture to the atlas
+        # NOTE: This is also called when re-building the atlas meaning we
+        #       need to support updating the texture coordinates for existing textures
         existing_slot = self._texture_uv_slots.get(texture.atlas_name)
         slot = existing_slot if existing_slot is not None else self._texture_uv_slots_free.popleft()
         self._texture_uv_slots[texture.atlas_name] = slot
+
         image_region = self.get_image_region_info(texture.image_data.hash)
         texture_region = copy.deepcopy(image_region)
-        # Since we copy the original image region we can always apply the
-        # transform without worrying about multiple transforms.
         texture_region.texture_coordinates = Transform.transform_texture_coordinates_order(
             texture_region.texture_coordinates, texture._vertex_order
         )
@@ -499,7 +578,7 @@ def _allocate_texture(self, texture: "Texture") -> Tuple[int, AtlasRegion]:
             self._texture_uv_data[offset + i] = texture_region.texture_coordinates[i]
 
         self._texture_uv_data_changed = True
-        self._textures.add(texture)
+        self._textures[texture.atlas_name] = texture
 
         return slot, texture_region
 
@@ -629,29 +708,25 @@ def remove(self, texture: "Texture") -> None:
         Remove a texture from the atlas.
 
         This doesn't erase the pixel data from the atlas texture
-        itself, but leaves the area unclaimed.
+        itself, but leaves the area unclaimed. The area will be
+        reclaimed when the atlas is rebuilt.
 
         :param texture: The texture to remove
         """
-        self._textures.remove(texture)
+        del self._textures[texture.atlas_name]
         # Reclaim the texture uv slot
         del self._texture_regions[texture.atlas_name]
         slot = self._texture_uv_slots[texture.atlas_name]
         del self._texture_uv_slots[texture.atlas_name]
         self._texture_uv_slots_free.appendleft(slot)
 
-        # Decrement the reference count for the image
-        self._image_ref_count.dec_ref(texture.image_data)
-        # print("Dec ref", texture.image_data.hash, self._image_ref_count.get_refs(texture.image_data))
-
         # Reclaim the image in the atlas if it's not used by any other texture
-        if self._image_ref_count.get_refs(texture.image_data) == 0:
+        if self._image_ref_count.dec_ref(texture.image_data) == 0:
             self._images.remove(texture.image_data)
             del self._image_regions[texture.image_data.hash]
             slot = self._image_uv_slots[texture.image_data.hash]
             del self._image_uv_slots[texture.image_data.hash]
             self._image_uv_slots_free.appendleft(slot)
-            # print("Reclaimed image", texture.image_data.hash)
 
     def update_texture_image(self, texture: "Texture"):
         """
@@ -714,7 +789,7 @@ def get_image_id(self, hash: str) -> int:
 
     def has_texture(self, texture: "Texture") -> bool:
         """Check if a texture is already in the atlas"""
-        return texture in self._textures
+        return texture.atlas_name in self._textures
 
     def has_image(self, image_data: "ImageData") -> bool:
         """Check if a image is already in the atlas"""
@@ -760,7 +835,7 @@ def resize(self, size: Tuple[int, int]) -> None:
 
         # Store old images and textures before clearing the atlas
         images = list(self._images)
-        textures = list(self._textures)
+        textures = self.textures
         # Clear the atlas without wiping the image and texture ids
         self.clear(clear_texture_ids=False, clear_image_ids=False, texture=False)
         for image in sorted(images, key=lambda x: x.height):
@@ -799,16 +874,17 @@ def resize(self, size: Tuple[int, int]) -> None:
 
         duration = time.perf_counter() - resize_start
         LOG.info("[%s] Atlas resize took %s seconds", id(self), duration)
-        # print(duration)
 
     def rebuild(self) -> None:
-        """Rebuild the underlying atlas texture.
+        """
+        Rebuild the underlying atlas texture.
 
         This method also tries to organize the textures more efficiently ordering them by size.
         The texture ids will persist so the sprite list don't need to be rebuilt.
         """
         # Hold a reference to the old textures
-        textures = list(self._textures)
+        textures = self.textures
+        self._image_ref_count = ImageDataRefCounter()
         # Clear the atlas but keep the uv slot mapping
         self.clear(clear_image_ids=False, clear_texture_ids=False)
         # Add textures back sorted by height to potentially make more room
@@ -833,13 +909,13 @@ def clear(
         """
         if texture:
             self._fbo.clear()
-        self._textures = WeakSet()
+        self._textures = WeakValueDictionary()
         self._images = WeakSet()
-        self._image_ref_count = ImageDataRefCounter()
         self._image_regions = dict()
         self._texture_regions = dict()
         self._allocator = Allocator(*self._size)
         if clear_image_ids:
+            self._image_ref_count = ImageDataRefCounter()
             self._image_uv_slots_free = deque(i for i in range(self._num_image_slots))
             self._image_uv_slots = dict()
         if clear_texture_ids:
@@ -1104,7 +1180,7 @@ def _check_size(self, size: Tuple[int, int]) -> None:
     def print_contents(self):
         """Debug method to print the contents of the atlas"""
         print("Textures:")
-        for texture in self._textures:
+        for texture in self.textures:
             print("->", texture)
         print("Images:")
         for image in self._images: