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

1from typing import Union 

2from PIL import Image, ImageDraw 

3import shapely 

4 

5 

6def draw_geometry( 

7 image_size: tuple[int, int], 

8 drawing_context: ImageDraw.ImageDraw, 

9 geometry: shapely.Geometry, 

10 value: int, 

11): 

12 """ 

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

14 

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

16 :param drawing_context: the drawing context to use 

17 :param geometry: the geometry to draw 

18 :param value: the value to use when drawing 

19 """ 

20 if isinstance(geometry, shapely.Point): 

21 _draw_point(drawing_context, geometry, value) 

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

23 _draw_line(drawing_context, geometry, value) 

24 elif isinstance(geometry, shapely.Polygon): 

25 _draw_polygon(image_size, drawing_context, geometry, value) 

26 elif isinstance(geometry, shapely.MultiPoint): 

27 for point in geometry.geoms: 

28 _draw_point(drawing_context, point, value) 

29 elif isinstance(geometry, shapely.MultiLineString): 

30 for line in geometry.geoms: 

31 _draw_line(drawing_context, line, value) 

32 elif isinstance(geometry, shapely.MultiPolygon): 

33 for polygon in geometry.geoms: 

34 _draw_polygon(image_size, drawing_context, polygon, value) 

35 elif isinstance(geometry, shapely.GeometryCollection): 

36 for g in geometry.geoms: 

37 draw_geometry(image_size, drawing_context, g, value) 

38 

39 

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

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

42 

43 

44def _draw_line( 

45 drawing_context: ImageDraw.ImageDraw, 

46 line: Union[shapely.LineString, shapely.LinearRing], 

47 value: int, 

48): 

49 drawing_context.line(line.coords, value) 

50 

51 

52def _draw_polygon( 

53 image_size: tuple[int, int], 

54 drawing_context: ImageDraw.ImageDraw, 

55 polygon: shapely.Polygon, 

56 value: int, 

57): 

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

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

60 if len(polygon.interiors) > 0: 

61 bitmap = Image.new("1", image_size) 

62 bitmap_draw = ImageDraw.Draw(bitmap) 

63 

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

65 for interior in polygon.interiors: 

66 bitmap_draw.polygon(interior.coords, 0) 

67 

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

69 else: 

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