Coverage for collectra/gui/segmentation.py: 92.16%

1"""Validate and atomically persist circular mask brush strokes.""" 

2 

3import math 

4import os 

5import tempfile 

6from pathlib import Path 

7 

8import numpy as np 

9from PIL import Image, ImageDraw 

10 

11 

12def paint_mask(item, strokes, orientation): 

13 from collectra.types.images import ImageSegmentation 

14 

15 if not isinstance(item, ImageSegmentation): 

16 raise ValueError("Only ImageSegmentation masks can be painted") 

17 if orientation != item.orientation.to_string(): 

18 raise ValueError( 

19 "The image orientation changed; reopen the segmentation before painting" 

20 ) 

21 if not isinstance(strokes, list) or not strokes: 

22 raise ValueError("No mask strokes supplied") 

23 pixels = item._mask_image().point(lambda value: 1 if value else 0) 

24 pixels = pixels.rotate(item.orientation.to_degree(), expand=True) 

25 before = np.asarray(pixels).copy() 

26 draw = ImageDraw.Draw(pixels) 

27 for stroke in strokes: 

28 if stroke.get("tool") not in {"brush", "eraser"}: 

29 raise ValueError("Unknown mask tool") 

30 radius = stroke.get("radius") 

31 if ( 

32 type(radius) not in (float, int) 

33 or not math.isfinite(radius) 

34 or not 0.5 <= radius <= 4096 

35 ): 

36 raise ValueError("Invalid brush radius") 

37 points = stroke.get("points") 

38 if not isinstance(points, list) or not points or len(points) > 100000: 

39 raise ValueError("Invalid brush points") 

40 for point in points: 

41 if ( 

42 not isinstance(point, list) 

43 or len(point) != 2 

44 or any( 

45 type(v) not in (int, float) or not math.isfinite(v) for v in point 

46 ) 

47 ): 

48 raise ValueError("Invalid brush point") 

49 if ( 

50 not -radius <= point[0] <= pixels.width + radius 

51 or not -radius <= point[1] <= pixels.height + radius 

52 ): 

53 raise ValueError("Brush point is outside the source image") 

54 fill = 1 if stroke["tool"] == "brush" else 0 

55 if len(points) > 1: 

56 draw.line( 

57 [tuple(p) for p in points], fill=fill, width=max(1, round(radius * 2)) 

58 ) 

59 for x, y in points: 

60 draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=fill) 

61 if pixels.getbbox() is None: 

62 raise ValueError( 

63 "A segmentation mask cannot be empty; leave at least one pixel" 

64 ) 

65 if np.array_equal(before, np.asarray(pixels)): 

66 return False 

67 pixels = pixels.rotate(-item.orientation.to_degree(), expand=True) 

68 path = Path(item._mask_path) 

69 descriptor, temporary = tempfile.mkstemp( 

70 prefix=f".{path.stem}-", suffix=".png", dir=path.parent 

71 ) 

72 try: 

73 with os.fdopen(descriptor, "wb") as stream: 

74 pixels.save(stream, format="PNG") 

75 os.chmod(temporary, path.stat().st_mode & 0o777) 

76 os.replace(temporary, path) 

77 finally: 

78 Path(temporary).unlink(missing_ok=True) 

79 return True