Path.py 2.6 KB

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