Chn11/MeshFormer

1import cv22import numpy as np3from PIL import Image4 5# marker choices6COLORS = [(246, 195, 203), (112, 221, 208)]7MARKERS = [1, 5]8 9def blend_seg(input_image, segmented_image, dot_locations, dot_labels, contour_color=(63, 126, 174)):10    input_image = np.array(input_image)11    segmented_image = np.array(segmented_image)12    13    # Create a mask for the foreground (non-transparent) pixels14    foreground_mask = segmented_image[:, :, 3] > 015 16    # Create a mask for the background (transparent) pixels17    background_mask = ~foreground_mask18 19    # Darken the background pixels20    darkened_background = input_image.copy()21    darkened_background[background_mask] = darkened_background[background_mask] * 0.52  # Adjust the multiplier as needed to control darkness22 23    # Create an empty mask for the boundary24    boundary_mask = np.zeros_like(segmented_image[:, :, 3], dtype=np.uint8)25    solid_boundary_mask = np.zeros_like(segmented_image[:, :, 3], dtype=np.uint8)26 27    # Find the contour of the segmented region28    contours, _ = cv2.findContours(29        segmented_image[:, :, 3], cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE30    )31 32    # Draw the boundary on the boundary mask33    cv2.drawContours(boundary_mask, contours, -1, (255), thickness=6)34    cv2.drawContours(solid_boundary_mask, contours, -1, (255), thickness=2)35 36    blur_mask = cv2.GaussianBlur(boundary_mask, (0, 0), sigmaX=4)37 38    # Create a mask for the contour region39    contour_region_mask = (blur_mask > 0)40 41    # Blend the contour color with the existing pixel colors42    result_image = darkened_background.copy()43    mask_weight = 0.9 * blur_mask[contour_region_mask, None]/25544    result_image[contour_region_mask] = (45        darkened_background[contour_region_mask] * (1-mask_weight) + np.array(contour_color) * mask_weight46    ).astype(np.uint8)47 48    # Overlay the contour on the result image without blending49    result_image[solid_boundary_mask > 0] = contour_color  # Set contour pixels to blue50 51    # Draw dots at the specified locations52    dot_radius = 653    for location, label in zip(dot_locations, dot_labels):54        if label:55            cv2.circle(result_image, location, dot_radius, COLORS[label], -1)56        else:57            cv2.drawMarker(result_image, location, COLORS[label], markerType=MARKERS[label],58                        markerSize=6, thickness=3)59 60    return Image.fromarray(result_image)61 62 63def blend_seg_pure(input_image, segmented_image, dot_locations, dot_labels):64    input_image = np.array(input_image)65    segmented_image = np.array(segmented_image)66    67    # Create a mask for the foreground (non-transparent) pixels68    foreground_mask = segmented_image[:, :, 3] > 069 70    # Blend the foreground71    red_foreground = input_image.copy()72    blend_weight = 0.873    red_foreground[foreground_mask] = red_foreground[foreground_mask] *(1-blend_weight) + np.array((255,0,0)) * blend_weight74 75    result_image = red_foreground76 77    # Draw dots at the specified locations78    dot_radius = 679    for location, label in zip(dot_locations, dot_labels):80        if label:81            cv2.circle(result_image, location, dot_radius, COLORS[label], -1)82        else:83            cv2.drawMarker(result_image, location, COLORS[label], markerType=MARKERS[label],84                        markerSize=6, thickness=3)85 86    return Image.fromarray(result_image)87