Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| # marker choices | |
| COLORS = [(246, 195, 203), (112, 221, 208)] | |
| MARKERS = [1, 5] | |
| def blend_seg(input_image, segmented_image, dot_locations, dot_labels, contour_color=(63, 126, 174)): | |
| input_image = np.array(input_image) | |
| segmented_image = np.array(segmented_image) | |
| # Create a mask for the foreground (non-transparent) pixels | |
| foreground_mask = segmented_image[:, :, 3] > 0 | |
| # Create a mask for the background (transparent) pixels | |
| background_mask = ~foreground_mask | |
| # Darken the background pixels | |
| darkened_background = input_image.copy() | |
| darkened_background[background_mask] = darkened_background[background_mask] * 0.52 # Adjust the multiplier as needed to control darkness | |
| # Create an empty mask for the boundary | |
| boundary_mask = np.zeros_like(segmented_image[:, :, 3], dtype=np.uint8) | |
| solid_boundary_mask = np.zeros_like(segmented_image[:, :, 3], dtype=np.uint8) | |
| # Find the contour of the segmented region | |
| contours, _ = cv2.findContours( | |
| segmented_image[:, :, 3], cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE | |
| ) | |
| # Draw the boundary on the boundary mask | |
| cv2.drawContours(boundary_mask, contours, -1, (255), thickness=6) | |
| cv2.drawContours(solid_boundary_mask, contours, -1, (255), thickness=2) | |
| blur_mask = cv2.GaussianBlur(boundary_mask, (0, 0), sigmaX=4) | |
| # Create a mask for the contour region | |
| contour_region_mask = (blur_mask > 0) | |
| # Blend the contour color with the existing pixel colors | |
| result_image = darkened_background.copy() | |
| mask_weight = 0.9 * blur_mask[contour_region_mask, None]/255 | |
| result_image[contour_region_mask] = ( | |
| darkened_background[contour_region_mask] * (1-mask_weight) + np.array(contour_color) * mask_weight | |
| ).astype(np.uint8) | |
| # Overlay the contour on the result image without blending | |
| result_image[solid_boundary_mask > 0] = contour_color # Set contour pixels to blue | |
| # Draw dots at the specified locations | |
| dot_radius = 6 | |
| for location, label in zip(dot_locations, dot_labels): | |
| if label: | |
| cv2.circle(result_image, location, dot_radius, COLORS[label], -1) | |
| else: | |
| cv2.drawMarker(result_image, location, COLORS[label], markerType=MARKERS[label], | |
| markerSize=6, thickness=3) | |
| return Image.fromarray(result_image) | |
| def blend_seg_pure(input_image, segmented_image, dot_locations, dot_labels): | |
| input_image = np.array(input_image) | |
| segmented_image = np.array(segmented_image) | |
| # Create a mask for the foreground (non-transparent) pixels | |
| foreground_mask = segmented_image[:, :, 3] > 0 | |
| # Blend the foreground | |
| red_foreground = input_image.copy() | |
| blend_weight = 0.8 | |
| red_foreground[foreground_mask] = red_foreground[foreground_mask] *(1-blend_weight) + np.array((255,0,0)) * blend_weight | |
| result_image = red_foreground | |
| # Draw dots at the specified locations | |
| dot_radius = 6 | |
| for location, label in zip(dot_locations, dot_labels): | |
| if label: | |
| cv2.circle(result_image, location, dot_radius, COLORS[label], -1) | |
| else: | |
| cv2.drawMarker(result_image, location, COLORS[label], markerType=MARKERS[label], | |
| markerSize=6, thickness=3) | |
| return Image.fromarray(result_image) | |