masking.py 2.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. from PIL import Image, ImageFilter, ImageOps
  2. def get_crop_region(mask, pad=0):
  3. """finds a rectangular region that contains all masked ares in an image. Returns (x1, y1, x2, y2) coordinates of the rectangle.
  4. For example, if a user has painted the top-right part of a 512x512 image, the result may be (256, 0, 512, 256)"""
  5. mask_img = mask if isinstance(mask, Image.Image) else Image.fromarray(mask)
  6. box = mask_img.getbbox()
  7. if box:
  8. x1, y1, x2, y2 = box
  9. else: # when no box is found
  10. x1, y1 = mask_img.size
  11. x2 = y2 = 0
  12. return max(x1 - pad, 0), max(y1 - pad, 0), min(x2 + pad, mask_img.size[0]), min(y2 + pad, mask_img.size[1])
  13. def expand_crop_region(crop_region, processing_width, processing_height, image_width, image_height):
  14. """expands crop region get_crop_region() to match the ratio of the image the region will processed in; returns expanded region
  15. for example, if user drew mask in a 128x32 region, and the dimensions for processing are 512x512, the region will be expanded to 128x128."""
  16. x1, y1, x2, y2 = crop_region
  17. ratio_crop_region = (x2 - x1) / (y2 - y1)
  18. ratio_processing = processing_width / processing_height
  19. if ratio_crop_region > ratio_processing:
  20. desired_height = (x2 - x1) / ratio_processing
  21. desired_height_diff = int(desired_height - (y2-y1))
  22. y1 -= desired_height_diff//2
  23. y2 += desired_height_diff - desired_height_diff//2
  24. if y2 >= image_height:
  25. diff = y2 - image_height
  26. y2 -= diff
  27. y1 -= diff
  28. if y1 < 0:
  29. y2 -= y1
  30. y1 -= y1
  31. if y2 >= image_height:
  32. y2 = image_height
  33. else:
  34. desired_width = (y2 - y1) * ratio_processing
  35. desired_width_diff = int(desired_width - (x2-x1))
  36. x1 -= desired_width_diff//2
  37. x2 += desired_width_diff - desired_width_diff//2
  38. if x2 >= image_width:
  39. diff = x2 - image_width
  40. x2 -= diff
  41. x1 -= diff
  42. if x1 < 0:
  43. x2 -= x1
  44. x1 -= x1
  45. if x2 >= image_width:
  46. x2 = image_width
  47. return x1, y1, x2, y2
  48. def fill(image, mask):
  49. """fills masked regions with colors from image using blur. Not extremely effective."""
  50. image_mod = Image.new('RGBA', (image.width, image.height))
  51. image_masked = Image.new('RGBa', (image.width, image.height))
  52. image_masked.paste(image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert('L')))
  53. image_masked = image_masked.convert('RGBa')
  54. for radius, repeats in [(256, 1), (64, 1), (16, 2), (4, 4), (2, 2), (0, 1)]:
  55. blurred = image_masked.filter(ImageFilter.GaussianBlur(radius)).convert('RGBA')
  56. for _ in range(repeats):
  57. image_mod.alpha_composite(blurred)
  58. return image_mod.convert("RGB")