Coverage for qubalab/objects/draw.py: 100%

1from typing import Union 

2from PIL import Image, ImageDraw 

3import shapely 

4 

5 

6def draw_geometry(image_size: tuple[int, int], drawing_context: ImageDraw.Draw, geometry: shapely.Geometry, value: int): 

7 """ 

8 Draw the provided Shapely geometry with the provided drawing context using the provided value. 

9 

10 :param image_size: the size of the image to draw on 

11 :param drawing_context: the drawing context to use 

12 :param geometry: the geometry to draw 

13 :param value: the value to use when drawing 

14 """ 

15 if (isinstance(geometry, shapely.Point)): 

16 _draw_point(drawing_context, geometry, value) 

17 elif (isinstance(geometry, Union[shapely.LineString, shapely.LinearRing])): 

18 _draw_line(drawing_context, geometry, value) 

19 elif (isinstance(geometry, shapely.Polygon)): 

20 _draw_polygon(image_size, drawing_context, geometry, value) 

21 elif (isinstance(geometry, shapely.MultiPoint)): 

22 for point in geometry.geoms: 

23 _draw_point(drawing_context, point, value) 

24 elif (isinstance(geometry, shapely.MultiLineString)): 

25 for line in geometry.geoms: 

26 _draw_line(drawing_context, line, value) 

27 elif (isinstance(geometry, shapely.MultiPolygon)): 

28 for polygon in geometry.geoms: 

29 _draw_polygon(image_size, drawing_context, polygon, value) 

30 elif (isinstance(geometry, shapely.GeometryCollection)): 

31 for g in geometry.geoms: 

32 draw_geometry(image_size, drawing_context, g, value) 

33 

34def _draw_point(drawing_context: ImageDraw, point: shapely.Point, value: int): 

35 drawing_context.point([point.x, point.y], value) 

36 

37def _draw_line(drawing_context: ImageDraw, line: Union[shapely.LineString, shapely.LinearRing], value: int): 

38 drawing_context.line(line.coords, value) 

39 

40def _draw_polygon(image_size: tuple[int, int], drawing_context: ImageDraw, polygon: shapely.Polygon, value: int): 

41 # If we have holes, we risk overpainting existing pixels with the background color 

42 # For that reason we need to create a separate binary image and then copy the values 

43 if len(polygon.interiors) > 0: 

44 bitmap = Image.new('1', image_size) 

45 bitmap_draw = ImageDraw.Draw(bitmap) 

46 

47 bitmap_draw.polygon(polygon.exterior.coords, 1) 

48 for interior in polygon.interiors: 

49 bitmap_draw.polygon(interior.coords, 0) 

50 

51 drawing_context.bitmap((0, 0), bitmap, value) 

52 else: 

53 drawing_context.polygon(polygon.exterior.coords, value)