(image, masks, mask_prompts, output_path, font_size=35, use_random_colors=False)
| 2 | import random |
| 3 | |
| 4 | def visualize_masks(image, masks, mask_prompts, output_path, font_size=35, use_random_colors=False): |
| 5 | # Create a blank image for overlays |
| 6 | overlay = Image.new('RGBA', image.size, (0, 0, 0, 0)) |
| 7 | |
| 8 | colors = [ |
| 9 | (165, 238, 173, 80), |
| 10 | (76, 102, 221, 80), |
| 11 | (221, 160, 77, 80), |
| 12 | (204, 93, 71, 80), |
| 13 | (145, 187, 149, 80), |
| 14 | (134, 141, 172, 80), |
| 15 | (157, 137, 109, 80), |
| 16 | (153, 104, 95, 80), |
| 17 | (165, 238, 173, 80), |
| 18 | (76, 102, 221, 80), |
| 19 | (221, 160, 77, 80), |
| 20 | (204, 93, 71, 80), |
| 21 | (145, 187, 149, 80), |
| 22 | (134, 141, 172, 80), |
| 23 | (157, 137, 109, 80), |
| 24 | (153, 104, 95, 80), |
| 25 | ] |
| 26 | # Generate random colors for each mask |
| 27 | if use_random_colors: |
| 28 | colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 80) for _ in range(len(masks))] |
| 29 | |
| 30 | # Font settings |
| 31 | try: |
| 32 | font = ImageFont.truetype("arial", font_size) # Adjust as needed |
| 33 | except IOError: |
| 34 | font = ImageFont.load_default(font_size) |
| 35 | |
| 36 | # Overlay each mask onto the overlay image |
| 37 | for mask, mask_prompt, color in zip(masks, mask_prompts, colors): |
| 38 | # Convert mask to RGBA mode |
| 39 | mask_rgba = mask.convert('RGBA') |
| 40 | mask_data = mask_rgba.getdata() |
| 41 | new_data = [(color if item[:3] == (255, 255, 255) else (0, 0, 0, 0)) for item in mask_data] |
| 42 | mask_rgba.putdata(new_data) |
| 43 | |
| 44 | # Draw the mask prompt text on the mask |
| 45 | draw = ImageDraw.Draw(mask_rgba) |
| 46 | mask_bbox = mask.getbbox() # Get the bounding box of the mask |
| 47 | text_position = (mask_bbox[0] + 10, mask_bbox[1] + 10) # Adjust text position based on mask position |
| 48 | draw.text(text_position, mask_prompt, fill=(255, 255, 255, 255), font=font) |
| 49 | |
| 50 | # Alpha composite the overlay with this mask |
| 51 | overlay = Image.alpha_composite(overlay, mask_rgba) |
| 52 | |
| 53 | # Composite the overlay onto the original image |
| 54 | result = Image.alpha_composite(image.convert('RGBA'), overlay) |
| 55 | |
| 56 | # Save or display the resulting image |
| 57 | result.save(output_path) |
| 58 | |
| 59 | return result |
no test coverage detected