Path.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. import cv2
  2. import numpy
  3. from gi.repository import Gdk
  4. class Path:
  5. def __init__(self, brush_size, brush_feather, scale, additive):
  6. self.points = []
  7. self.brush_size = brush_size
  8. self.scale = scale
  9. self.brush_feather = brush_feather
  10. self.additive = additive
  11. self.has_rendered = False
  12. def add_point(self, x, y, previewShape = None, fill = None):
  13. self.points.append([x,y])
  14. self.has_rendered = False
  15. if(previewShape != None and fill != None and len(self.points) > 1):
  16. preview = numpy.zeros(previewShape, dtype=numpy.uint8)
  17. points = self.points[-2:]
  18. sx = points[0][0]
  19. sy = points[0][1]
  20. fx = points[1][0]
  21. fy = points[1][1]
  22. cv2.line(preview, (sx, sy), (fx, fy), (fill), int(self.brush_size), cv2.LINE_4)
  23. return preview
  24. def get_mask_map(self, mask):
  25. self.has_rendered = True
  26. if(self.additive):
  27. return self.draw_additive_path(mask)
  28. else:
  29. return self.draw_subtractive_path(mask)
  30. def draw_subtractive_path(self, mask):
  31. self.draw_points(mask, self.points, 0)
  32. return mask
  33. def draw_additive_path(self, mask):
  34. map = numpy.zeros(mask.shape, dtype=numpy.uint8)
  35. self.draw_points(map, self.points, 255)
  36. if(self.brush_feather > 1):
  37. blur_size = 2 * round((round(self.brush_feather) + 1) / 2) - 1
  38. map = cv2.blur(map, (int(blur_size), int(blur_size)))
  39. map = map[:, :, numpy.newaxis]
  40. # Painful workaround for int overflow
  41. mask2 = mask.astype(numpy.uint16)
  42. mask2 = mask2 + map
  43. mask2[mask2 > 255] = 255
  44. mask = mask2.astype(numpy.uint8)
  45. return mask
  46. def draw_points(self, mask, points, value):
  47. for i in range(len(points) - 2):
  48. sx = int(self.points[i][0] * self.scale)
  49. sy = int(self.points[i][1] * self.scale)
  50. fx = int(self.points[i + 1][0] * self.scale)
  51. fy = int(self.points[i + 1][1] * self.scale)
  52. # print("[PATH.PY]", sx, sy, fx, fy)
  53. cv2.line(mask, (sx, sy), (fx, fy), (value), int(self.brush_size * self.scale), cv2.LINE_AA)
  54. return mask
  55. def get_path_dict(self):
  56. return {
  57. "brush_size": self.brush_size,
  58. "scale": self.scale,
  59. "brush_feather": self.brush_feather,
  60. "additive": self.additive,
  61. "points": self.points
  62. }
  63. @staticmethod
  64. def new_from_dict(dict):
  65. path = Path(dict["brush_size"], dict["brush_feather"], dict["scale"], dict["additive"])
  66. path.points = dict["points"]
  67. return path