Skip to content

Commit

Permalink
Add an OpenCV backend. (#103)
Browse files Browse the repository at this point in the history
* Add an OpenCV backend.

On conda-forge because all Pillow and OpenCV are built with with
jpegturbo i can't really mesure any difference with the 3 different jpeg
encoders (simplejpeg installed from pip)

* Update

* importorskip
  • Loading branch information
hmaarrfk authored Nov 5, 2024
1 parent 8556297 commit 06b2690
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 1 deletion.
2 changes: 1 addition & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ jobs:
shell: bash
run: |
python -m pip install --upgrade pip
pip install -U jupyter_packaging pytest simplejpeg pillow
pip install -U jupyter_packaging pytest simplejpeg pillow opencv-python-headless
pip install .
rm -rf ./jupyter_rfb ./build ./egg-info
- name: Test with pytest
Expand Down
23 changes: 23 additions & 0 deletions jupyter_rfb/_jpg.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,12 +99,35 @@ def _encode(self, array, quality):
return f.getvalue()


class OpenCVJpegEncoder(JpegEncoder):
"""A JPEG encoder using the OpenCV library."""

def __init__(self):
import cv2

self.cv2 = cv2

def _encode(self, array, quality):
if len(array.shape) == 3 and array.shape[2] == 3:
# Convert RGB to BGR if needed (assume input is RGB)
array = self.cv2.cvtColor(array, self.cv2.COLOR_RGB2BGR)

# Encode with the specified quality
encode_param = [self.cv2.IMWRITE_JPEG_QUALITY, quality]
success, encoded_image = self.cv2.imencode(".jpg", array, encode_param)
if not success:
raise RuntimeError("OpenCV failed to encode image")

return encoded_image.tobytes()


def select_encoder():
"""Select an encoder."""

for cls in [
SimpleJpegEncoder, # simplejpeg is fast and lean
PillowJpegEncoder, # pillow is commonly available
OpenCVJpegEncoder, # opencv is readily installed in conda environments
]:
try:
return cls()
Expand Down
15 changes: 15 additions & 0 deletions tests/test_jpg.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
"""Test jpg module."""

import numpy as np
import pytest
from pytest import raises

from jupyter_rfb._jpg import (
array2jpg,
select_encoder,
SimpleJpegEncoder,
PillowJpegEncoder,
OpenCVJpegEncoder,
)


Expand All @@ -28,18 +30,28 @@ def test_array2jpg():

def test_simplejpeg_jpeg_encoder():
"""Test the simplejpeg encoder."""
pytest.importorskip("simplejpeg")
encoder = SimpleJpegEncoder()
_perform_checks(encoder)
_perform_error_checks(encoder)


def test_pillow_jpeg_encoder():
"""Test the pillow encoder."""
pytest.importorskip("PIL")
encoder = PillowJpegEncoder()
_perform_checks(encoder)
_perform_error_checks(encoder)


def test_opencv_jpeg_encoder():
"""Test the opencv encoder."""
pytest.importorskip("cv2")
encoder = OpenCVJpegEncoder()
_perform_checks(encoder)
_perform_error_checks(encoder)


def _perform_checks(encoder):
# RGB
im = get_random_im(100, 100, 3)
Expand Down Expand Up @@ -111,9 +123,11 @@ def test_select_encoder():
# Sabotage
simple_init = SimpleJpegEncoder.__init__
pillow_init = PillowJpegEncoder.__init__
cv2_init = OpenCVJpegEncoder.__init__
try:
SimpleJpegEncoder.__init__ = lambda self: raise_importerror()
PillowJpegEncoder.__init__ = lambda self: raise_importerror()
OpenCVJpegEncoder.__init__ = lambda self: raise_importerror()

encoder = select_encoder()
assert not isinstance(encoder, (SimpleJpegEncoder, PillowJpegEncoder))
Expand All @@ -125,3 +139,4 @@ def test_select_encoder():
finally:
SimpleJpegEncoder.__init__ = simple_init
PillowJpegEncoder.__init__ = pillow_init
OpenCVJpegEncoder.__init__ = cv2_init

0 comments on commit 06b2690

Please sign in to comment.