Building the ImageManager Class in Django

Introduction to the Image Manager

Welcome to the second lesson of our course on building an image generation service with Django . In our previous lesson, we created the PromptManager class to handle the formatting of prompts for our image generation process. Now, we'll build upon that foundation by creating the ImageManager class, which will be responsible for handling the images once they're generated. The ImageManager serves as a crucial component in our application architecture. While the PromptManager prepares the instructions for image generation, the ImageManager takes care of what happens after an image is created. Its primary responsibilities include: Storing generated images along with their associated prompts Converting images to a web-friendly format (base64) Providing access to the collection of stored images This component is essential because web applications can't directly work with raw image data. By converting images to base64 strings, we can easily embed them in HTML or send them as JSON responses in our API. Additionally, keeping track of which prompt generated which image allows us to maintain a history of generations and potentially reuse successful prompts. Let's dive into building this important piece of our image generation service.

Setting Up the Image Manager Class

To begin, we'll create a new file called image_manager.py in our app/models directory. This file will contain our ImageManager class. First, let's set up the basic structure and understand the imports we'll need: Python import base64 from PIL import Image from io import BytesIO class ImageManager: def __init__(self): self.images = [] Let's break down these imports: base64 : This module provides functions for encoding binary data to ASCII characters and decoding such encodings back to binary data. We'll use it to convert our images to a string format that can be easily transmitted over HTTP. PIL (Python Imaging Library): We import the Image class from PIL, which allows us to open, manipulate, and save many different image file formats. In our case, we'll use it to process the image data. BytesIO : This is a class from the io module that implements a file-like interface for reading and writing bytes in memory rather than to a disk file. It's useful for handling binary data like images. Our ImageManager class starts with a simple constructor that initializes an empty list called images . This list will store dictionaries containing information about each image, including: A unique identifier The prompt that generated the image The base64 -encoded image data This simple data structure allows us to keep track of all generated images and their associated metadata in memory. Keep in mind that in-memory storage is volatile — data will be lost every time the server restarts. It's suitable for prototyping and local testing, but in real-world applications, you would typically store this information in a database instead.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Converting Images to Base64 Format

One of the key functions of our ImageManager class is to convert image data into a format that can be easily used in web applications. Base64 encoding is perfect for this purpose as it transforms binary data into a string of ASCII characters that can be included directly in HTML or JSON.

For example, you can embed a base64-encoded image directly into an HTML <img> tag using the src="data:image/jpeg;base64,..." format. This avoids having to serve the image as a separate static file, which simplifies delivery for small-scale applications or APIs that return dynamic image content. However, note that this technique is inefficient for large images or high-volume applications due to increased payload size and rendering time.

Let's implement the image_to_base64 method:

Python

def image_to_base64(self, image_data):
    try:
        image_bytes = image_data if isinstance(image_data, bytes) else image_data.image.image_bytes
        image = Image.open(BytesIO(image_bytes))
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        return base64.b64encode(buffered.getvalue()).decode("utf-8")
    except Exception as e:
        raise RuntimeError(f"Error processing image: {str(e)}")

def image_to_base64(self, image_data):
    try:
        image_bytes = image_data if isinstance(image_data, bytes) else image_data.image.image_bytes
        image = Image.open(BytesIO(image_bytes))
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        return base64.b64encode(buffered.getvalue()).decode("utf-8")
    except Exception as e:
        raise RuntimeError(f"Error processing image: {str(e)}")

This method takes image_data as input, which will be the response from our image generation API. Let's walk through what's happening:

We extract the raw image bytes from the API response structure (the raw image_data bytes).
We use BytesIO to create a file-like object from these bytes, which PIL's Image.open() can read.
We create another BytesIO buffer to hold the processed image.
We save the image to this buffer in JPEG format.
We encode the buffer's contents as base64 and then decode it to a UTF-8 string (rather than leaving it as bytes).
If any errors occur during this process, we catch them and raise a more informative RuntimeError.

The error handling is particularly important here because image processing can fail for various reasons, such as corrupted image data or memory issues. By catching exceptions and providing clear error messages, we make debugging much easier.

Testing the Image Manager

Now that we've implemented our ImageManager class, let's create a simple test script to verify that it works correctly. We'll create this in our app/tests.py file:

Python

from django.test import TestCase
from app.models.image_manager import ImageManager

class ImageManagerTestCase(TestCase):
    def setUp(self):
        self.image_manager = ImageManager()
        self.fake_image_data = b"fake image data"  # Simulated image data (use real data when integrating)
        self.prompt = "Test prompt for ImageManager"

    def test_add_image(self):
        result = self.image_manager.add_image(self.prompt, self.fake_image_data)
        self.assertIsNotNone(result)
        self.assertEqual(len(self.image_manager.get_images()), 1)

    def test_get_images(self):
        self.image_manager.add_image(self.prompt, self.fake_image_data)
        images = self.image_manager.get_images()
        self.assertEqual(len(images), 1)
        self.assertEqual(images[0]['prompt'], self.prompt)

from django.test import TestCase
from app.models.image_manager import ImageManager

class ImageManagerTestCase(TestCase):
    def setUp(self):
        self.image_manager = ImageManager()
        self.fake_image_data = b"fake image data"  # Simulated image data (use real data when integrating)
        self.prompt = "Test prompt for ImageManager"

    def test_add_image(self):
        result = self.image_manager.add_image(self.prompt, self.fake_image_data)
        self.assertIsNotNone(result)
        self.assertEqual(len(self.image_manager.get_images()), 1)

    def test_get_images(self):
        self.image_manager.add_image(self.prompt, self.fake_image_data)
        images = self.image_manager.get_images()
        self.assertEqual(len(images), 1)
        self.assertEqual(images[0]['prompt'], self.prompt)

In this test script, we:

Use Django's TestCase to create a test class for our ImageManager.
Set up an instance of the ImageManager and define some fake image data and a test prompt in the setUp method.
Implement a test for the add_image method to ensure it adds an image correctly.
Implement a test for the get_images method to ensure it retrieves stored images correctly.

Introduction to the Image Manager

Setting Up the Image Manager Class

Join the 1M+ learners on CodeSignal

Converting Images to Base64 Format

Building Image Collection Methods

The add_image Method

The get_images Method

Testing the Image Manager

Summary and Practice Preview