Skip to content
Snippets Groups Projects
Commit e06951ea authored by Erwan Rouchet's avatar Erwan Rouchet Committed by Bastien Abadie
Browse files

Support rotation and mirroring in open_image

parent 3e2d0b11
No related branches found
No related tags found
1 merge request!130Support rotation and mirroring in open_image
Pipeline #78767 passed
Stage: release
Hide whitespace changes
Inline Side-by-side
Showing
with 150 additions and 31 deletions
arkindex_worker/cache.py
+ 9
− 1
View file @ e06951ea
......@@ -55,6 +55,8 @@ class CachedElement(Model):
type = CharField(max_length=50)
image = ForeignKeyField(CachedImage, backref="elements", null=True)
polygon = JSONField(null=True)
rotation_angle = IntegerField(default=0)
mirrored = BooleanField(default=False)
initial = BooleanField(default=False)
worker_version_id = UUIDField(null=True)
......@@ -108,7 +110,13 @@ class CachedElement(Model):
if not url.endswith("/"):
url += "/"
return open_image(f"{url}{box}/{resize}/0/default.jpg", *args, **kwargs)
return open_image(
f"{url}{box}/{resize}/0/default.jpg",
*args,
rotation_angle=self.rotation_angle,
mirrored=self.mirrored,
**kwargs,
)
class CachedTranscription(Model):
......
arkindex_worker/image.py
+ 7
− 1
View file @ e06951ea
......@@ -21,7 +21,7 @@ DOWNLOAD_TIMEOUT = (30, 60)
BoundingBox = namedtuple("BoundingBox", ["x", "y", "width", "height"])
def open_image(path, mode="RGB"):
def open_image(path, mode="RGB", rotation_angle=0, mirrored=False):
"""
Open an image from a path or a URL
"""
......@@ -40,6 +40,12 @@ def open_image(path, mode="RGB"):
if image.mode != mode:
image = image.convert(mode)
if mirrored:
image = image.transpose(Image.FLIP_LEFT_RIGHT)
if rotation_angle:
image = image.rotate(-rotation_angle, expand=True)
return image
......
arkindex_worker/models.py
+ 12
− 4
View file @ e06951ea
......@@ -135,11 +135,19 @@ class Element(MagicDict):
else:
resize = "full"
if use_full_image:
url = self.image_url(resize)
else:
url = self.resize_zone_url(resize)
try:
if use_full_image:
return open_image(self.image_url(resize), *args, **kwargs)
else:
return open_image(self.resize_zone_url(resize), *args, **kwargs)
return open_image(
url,
*args,
rotation_angle=self.rotation_angle,
mirrored=self.mirrored,
**kwargs
)
except HTTPError as e:
if (
self.zone.image.get("s3_url") is not None
......
tests/data/mirrored_image.jpg 0 → 100644
+ 0
− 0
View file @ e06951ea
tests/data/mirrored_image.jpg

14.4 KiB

tests/data/rotated_image.jpg 0 → 100644
+ 0
− 0
View file @ e06951ea
tests/data/rotated_image.jpg

21.9 KiB

tests/data/rotated_mirrored_image.jpg 0 → 100644
+ 0
− 0
View file @ e06951ea
tests/data/rotated_mirrored_image.jpg

21.8 KiB

tests/test_cache.py
+ 4
− 2
View file @ e06951ea
......@@ -55,7 +55,7 @@ def test_create_tables(tmp_path):
create_tables()
expected_schema = """CREATE TABLE "classifications" ("id" TEXT NOT NULL PRIMARY KEY, "element_id" TEXT NOT NULL, "class_name" TEXT NOT NULL, "confidence" REAL NOT NULL, "state" VARCHAR(10) NOT NULL, "worker_version_id" TEXT NOT NULL, FOREIGN KEY ("element_id") REFERENCES "elements" ("id"))
CREATE TABLE "elements" ("id" TEXT NOT NULL PRIMARY KEY, "parent_id" TEXT, "type" VARCHAR(50) NOT NULL, "image_id" TEXT, "polygon" text, "initial" INTEGER NOT NULL, "worker_version_id" TEXT, FOREIGN KEY ("image_id") REFERENCES "images" ("id"))
CREATE TABLE "elements" ("id" TEXT NOT NULL PRIMARY KEY, "parent_id" TEXT, "type" VARCHAR(50) NOT NULL, "image_id" TEXT, "polygon" text, "rotation_angle" INTEGER NOT NULL, "mirrored" INTEGER NOT NULL, "initial" INTEGER NOT NULL, "worker_version_id" TEXT, FOREIGN KEY ("image_id") REFERENCES "images" ("id"))
CREATE TABLE "entities" ("id" TEXT NOT NULL PRIMARY KEY, "type" VARCHAR(50) NOT NULL, "name" TEXT NOT NULL, "validated" INTEGER NOT NULL, "metas" text, "worker_version_id" TEXT NOT NULL)
CREATE TABLE "images" ("id" TEXT NOT NULL PRIMARY KEY, "width" INTEGER NOT NULL, "height" INTEGER NOT NULL, "url" TEXT NOT NULL)
CREATE TABLE "transcription_entities" ("transcription_id" TEXT NOT NULL, "entity_id" TEXT NOT NULL, "offset" INTEGER NOT NULL CHECK (offset >= 0), "length" INTEGER NOT NULL CHECK (length > 0), "worker_version_id" TEXT NOT NULL, PRIMARY KEY ("transcription_id", "entity_id"), FOREIGN KEY ("transcription_id") REFERENCES "transcriptions" ("id"), FOREIGN KEY ("entity_id") REFERENCES "entities" ("id"))
......@@ -189,7 +189,9 @@ def test_element_open_image(
assert elt.open_image(max_size=max_size) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(expected_url)
assert open_mock.call_args == mocker.call(
expected_url, mirrored=False, rotation_angle=0
)
def test_element_open_image_requires_image():
......
tests/test_element.py
+ 99
− 23
View file @ e06951ea
......@@ -59,13 +59,17 @@ def test_open_image(mocker):
"server": {"max_width": None, "max_height": None},
},
"polygon": [[0, 0], [181, 0], [181, 240], [0, 240], [0, 0]],
}
}
},
"rotation_angle": 0,
"mirrored": False,
},
)
assert elt.open_image(use_full_image=True) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg"
"http://something/full/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
......@@ -83,26 +87,34 @@ def test_open_image_resize_portrait(mocker):
"server": {"max_width": None, "max_height": None},
},
"polygon": [[0, 0], [400, 0], [400, 600], [0, 600], [0, 0]],
}
},
"rotation_angle": 0,
"mirrored": False,
}
)
# Resize = original size
assert elt.open_image(max_size=600, use_full_image=True) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg"
"http://something/full/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
# Resize = smaller height
assert elt.open_image(max_size=400, use_full_image=True) == "an image!"
assert open_mock.call_count == 2
assert open_mock.call_args == mocker.call(
"http://something/full/266,400/0/default.jpg"
"http://something/full/266,400/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
# Resize = bigger height
assert elt.open_image(max_size=800, use_full_image=True) == "an image!"
assert open_mock.call_count == 3
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg"
"http://something/full/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
......@@ -120,13 +132,17 @@ def test_open_image_resize_partial_element(mocker):
"server": {"max_width": None, "max_height": None},
},
"polygon": [[0, 0], [200, 0], [200, 600], [0, 600], [0, 0]],
}
},
"rotation_angle": 0,
"mirrored": False,
}
)
assert elt.open_image(max_size=400, use_full_image=True) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg"
"http://something/full/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
......@@ -144,26 +160,34 @@ def test_open_image_resize_landscape(mocker):
"server": {"max_width": None, "max_height": None},
},
"polygon": [[0, 0], [600, 0], [600, 400], [0, 400], [0, 0]],
}
},
"rotation_angle": 0,
"mirrored": False,
}
)
# Resize = original size
assert elt.open_image(max_size=600, use_full_image=True) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg"
"http://something/full/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
# Resize = smaller width
assert elt.open_image(max_size=400, use_full_image=True) == "an image!"
assert open_mock.call_count == 2
assert open_mock.call_args == mocker.call(
"http://something/full/400,266/0/default.jpg"
"http://something/full/400,266/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
# Resize = bigger width
assert elt.open_image(max_size=800, use_full_image=True) == "an image!"
assert open_mock.call_count == 3
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg"
"http://something/full/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
......@@ -181,26 +205,34 @@ def test_open_image_resize_square(mocker):
"server": {"max_width": None, "max_height": None},
},
"polygon": [[0, 0], [400, 0], [400, 400], [0, 400], [0, 0]],
}
},
"rotation_angle": 0,
"mirrored": False,
}
)
# Resize = original size
assert elt.open_image(max_size=400, use_full_image=True) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg"
"http://something/full/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
# Resize = smaller
assert elt.open_image(max_size=200, use_full_image=True) == "an image!"
assert open_mock.call_count == 2
assert open_mock.call_args == mocker.call(
"http://something/full/200,200/0/default.jpg"
"http://something/full/200,200/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
# Resize = bigger
assert elt.open_image(max_size=800, use_full_image=True) == "an image!"
assert open_mock.call_count == 3
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg"
"http://something/full/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
......@@ -232,11 +264,17 @@ def test_open_image_s3(mocker):
"arkindex_worker.models.open_image", return_value="an image!"
)
elt = Element(
{"zone": {"image": {"url": "http://something", "s3_url": "http://s3url"}}}
{
"zone": {"image": {"url": "http://something", "s3_url": "http://s3url"}},
"rotation_angle": 0,
"mirrored": False,
}
)
assert elt.open_image(use_full_image=True) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call("http://s3url")
assert open_mock.call_args == mocker.call(
"http://s3url", rotation_angle=0, mirrored=False
)
def test_open_image_s3_retry(mocker):
......@@ -252,6 +290,8 @@ def test_open_image_s3_retry(mocker):
{
"id": "cafe",
"zone": {"image": {"url": "http://something", "s3_url": "http://oldurl"}},
"rotation_angle": 0,
"mirrored": False,
}
)
......@@ -270,6 +310,8 @@ def test_open_image_s3_retry_once(mocker):
{
"id": "cafe",
"zone": {"image": {"url": "http://something", "s3_url": "http://oldurl"}},
"rotation_angle": 0,
"mirrored": False,
}
)
......@@ -286,13 +328,17 @@ def test_open_image_use_full_image_false(mocker):
"zone": {
"image": {"url": "http://something", "s3_url": "http://s3url"},
"url": "http://zoneurl/0,0,400,600/full/0/default.jpg",
}
},
"rotation_angle": 0,
"mirrored": False,
}
)
assert elt.open_image(use_full_image=False) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(
"http://zoneurl/0,0,400,600/full/0/default.jpg"
"http://zoneurl/0,0,400,600/full/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
......@@ -311,14 +357,44 @@ def test_open_image_resize_use_full_image_false(mocker):
},
"polygon": [[0, 0], [400, 0], [400, 600], [0, 600], [0, 0]],
"url": "http://zoneurl/0,0,400,600/full/0/default.jpg",
}
},
"rotation_angle": 0,
"mirrored": False,
}
)
# Resize = smaller
assert elt.open_image(max_size=200, use_full_image=False) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(
"http://zoneurl/0,0,400,600/133,200/0/default.jpg"
"http://zoneurl/0,0,400,600/133,200/0/default.jpg",
rotation_angle=0,
mirrored=False,
)
def test_open_image_rotation_mirror(mocker):
open_mock = mocker.patch(
"arkindex_worker.models.open_image", return_value="an image!"
)
elt = Element(
{
"zone": {
"image": {
"url": "http://something",
"server": {"max_width": None, "max_height": None},
},
"polygon": [[0, 0], [181, 0], [181, 240], [0, 240], [0, 0]],
},
"rotation_angle": 42,
"mirrored": True,
},
)
assert elt.open_image(use_full_image=True) == "an image!"
assert open_mock.call_count == 1
assert open_mock.call_args == mocker.call(
"http://something/full/full/0/default.jpg",
rotation_angle=42,
mirrored=True,
)
......
tests/test_image.py
+ 19
− 0
View file @ e06951ea
......@@ -11,6 +11,9 @@ from arkindex_worker.image import download_tiles, open_image
FIXTURES = Path(__file__).absolute().parent / "data"
TILE = FIXTURES / "test_image.jpg"
FULL_IMAGE = FIXTURES / "tiled_image.jpg"
ROTATED_IMAGE = FIXTURES / "rotated_image.jpg"
MIRRORED_IMAGE = FIXTURES / "mirrored_image.jpg"
ROTATED_MIRRORED_IMAGE = FIXTURES / "rotated_mirrored_image.jpg"
def _root_mean_square(img_a, img_b):
......@@ -127,3 +130,19 @@ def test_open_image(path, is_local, monkeypatch):
assert image.size == (1, 10)
else:
assert image.size == (10, 1)
@pytest.mark.parametrize(
"rotation_angle,mirrored,expected_path",
(
(0, False, TILE),
(45, False, ROTATED_IMAGE),
(0, True, MIRRORED_IMAGE),
(45, True, ROTATED_MIRRORED_IMAGE),
),
)
def test_open_image_rotate_mirror(rotation_angle, mirrored, expected_path):
expected = Image.open(expected_path).convert("RGB")
actual = open_image(str(TILE), rotation_angle=rotation_angle, mirrored=mirrored)
actual.save(f"/tmp/{rotation_angle}_{mirrored}.jpg")
assert _root_mean_square(expected, actual) <= 15.0
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment