File size: 3,942 Bytes
14114e8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import cv2
import numpy as np
import pycocotools.mask as mask_utils
from PIL import Image
from .helpers.visualizer import Visualizer
from .helpers.zoom_in import render_zoom_in
def visualize(
input_json: dict,
zoom_in_index: int | None = None,
mask_alpha: float = 0.15,
label_mode: str = "1",
font_size_multiplier: float = 1.2,
boarder_width_multiplier: float = 0,
):
"""
Unified visualization function.
If zoom_in_index is None:
- Render all masks in input_json (equivalent to visualize_masks_from_result_json).
- Returns: PIL.Image
If zoom_in_index is provided:
- Returns two PIL.Images:
1) Output identical to zoom_in_and_visualize(input_json, index).
2) The same instance rendered via the general overlay using the color
returned by (1), equivalent to calling visualize_masks_from_result_json
on a single-mask json_i with color=color_hex.
"""
# Common fields
orig_h = int(input_json["orig_img_h"])
orig_w = int(input_json["orig_img_w"])
img_path = input_json["original_image_path"]
# ---------- Mode A: Full-scene render ----------
if zoom_in_index is None:
boxes = np.array(input_json["pred_boxes"])
rle_masks = [
{"size": (orig_h, orig_w), "counts": rle}
for rle in input_json["pred_masks"]
]
binary_masks = [mask_utils.decode(rle) for rle in rle_masks]
img_bgr = cv2.imread(img_path)
if img_bgr is None:
raise FileNotFoundError(f"Could not read image: {img_path}")
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
viz = Visualizer(
img_rgb,
font_size_multiplier=font_size_multiplier,
boarder_width_multiplier=boarder_width_multiplier,
)
viz.overlay_instances(
boxes=boxes,
masks=rle_masks,
binary_masks=binary_masks,
assigned_colors=None,
alpha=mask_alpha,
label_mode=label_mode,
)
pil_all_masks = Image.fromarray(viz.output.get_image())
return pil_all_masks
# ---------- Mode B: Zoom-in pair ----------
else:
idx = int(zoom_in_index)
num_masks = len(input_json.get("pred_masks", []))
if idx < 0 or idx >= num_masks:
raise ValueError(f"zoom_in_index {idx} is out of range (0..{num_masks-1}).")
# (1) Replicate zoom_in_and_visualize
object_data = {
"labels": [{"noun_phrase": f"mask_{idx}"}],
"segmentation": {
"counts": input_json["pred_masks"][idx],
"size": [orig_h, orig_w],
},
}
pil_img = Image.open(img_path)
pil_mask_i_zoomed, color_hex = render_zoom_in(
object_data, pil_img, mask_alpha=mask_alpha
)
# (2) Single-instance render with the same color
boxes_i = np.array([input_json["pred_boxes"][idx]])
rle_i = {"size": (orig_h, orig_w), "counts": input_json["pred_masks"][idx]}
bin_i = mask_utils.decode(rle_i)
img_bgr_i = cv2.imread(img_path)
if img_bgr_i is None:
raise FileNotFoundError(f"Could not read image: {img_path}")
img_rgb_i = cv2.cvtColor(img_bgr_i, cv2.COLOR_BGR2RGB)
viz_i = Visualizer(
img_rgb_i,
font_size_multiplier=font_size_multiplier,
boarder_width_multiplier=boarder_width_multiplier,
)
viz_i.overlay_instances(
boxes=boxes_i,
masks=[rle_i],
binary_masks=[bin_i],
assigned_colors=[color_hex],
alpha=mask_alpha,
label_mode=label_mode,
)
pil_mask_i = Image.fromarray(viz_i.output.get_image())
return pil_mask_i, pil_mask_i_zoomed
|