Mirrors/Track-Anything
1
0
Fork 0
mirror of https://github.com/gaomingqi/Track-Anything.git synced 2025-12-15 16:07:51 +01:00
Code Issues Packages Projects Releases Wiki Activity

add mask and point painter

Browse Source
This commit is contained in:
gaomingqi
2023-04-13 23:17:17 +08:00
parent 0618c12d63
commit 81009e4447
11 changed files with 193 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

@@ -2,4 +2,5 @@ __pycache__/
.vscode/
docs/
*.pth
*.mp4
debug_images/

BIN
images/mask_painter.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 701 KiB

BIN
images/painter_input_image.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 438 KiB

BIN
images/painter_input_mask.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 23 KiB

BIN
images/painter_output_image.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 704 KiB

BIN
images/painter_output_image__.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 4.9 KiB

BIN
images/point_painter.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 739 KiB

BIN
images/point_painter_1.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 739 KiB

BIN
images/point_painter_2.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 739 KiB

191
tools/painter.py Normal file
View File

@@ -0,0 +1,191 @@
# paint masks, contours, or points on images, with specified colors
import cv2
import torch
import numpy as np
from PIL import Image
import copy
import time
def colormap(rgb=True):
color_list = np.array(
[
0.000, 0.000, 0.000,
1.000, 1.000, 1.000,
1.000, 0.498, 0.313,
0.392, 0.581, 0.929,
0.000, 0.447, 0.741,
0.850, 0.325, 0.098,
0.929, 0.694, 0.125,
0.494, 0.184, 0.556,
0.466, 0.674, 0.188,
0.301, 0.745, 0.933,
0.635, 0.078, 0.184,
0.300, 0.300, 0.300,
0.600, 0.600, 0.600,
1.000, 0.000, 0.000,
1.000, 0.500, 0.000,
0.749, 0.749, 0.000,
0.000, 1.000, 0.000,
0.000, 0.000, 1.000,
0.667, 0.000, 1.000,
0.333, 0.333, 0.000,
0.333, 0.667, 0.000,
0.333, 1.000, 0.000,
0.667, 0.333, 0.000,
0.667, 0.667, 0.000,
0.667, 1.000, 0.000,
1.000, 0.333, 0.000,
1.000, 0.667, 0.000,
1.000, 1.000, 0.000,
0.000, 0.333, 0.500,
0.000, 0.667, 0.500,
0.000, 1.000, 0.500,
0.333, 0.000, 0.500,
0.333, 0.333, 0.500,
0.333, 0.667, 0.500,
0.333, 1.000, 0.500,
0.667, 0.000, 0.500,
0.667, 0.333, 0.500,
0.667, 0.667, 0.500,
0.667, 1.000, 0.500,
1.000, 0.000, 0.500,
1.000, 0.333, 0.500,
1.000, 0.667, 0.500,
1.000, 1.000, 0.500,
0.000, 0.333, 1.000,
0.000, 0.667, 1.000,
0.000, 1.000, 1.000,
0.333, 0.000, 1.000,
0.333, 0.333, 1.000,
0.333, 0.667, 1.000,
0.333, 1.000, 1.000,
0.667, 0.000, 1.000,
0.667, 0.333, 1.000,
0.667, 0.667, 1.000,
0.667, 1.000, 1.000,
1.000, 0.000, 1.000,
1.000, 0.333, 1.000,
1.000, 0.667, 1.000,
0.167, 0.000, 0.000,
0.333, 0.000, 0.000,
0.500, 0.000, 0.000,
0.667, 0.000, 0.000,
0.833, 0.000, 0.000,
1.000, 0.000, 0.000,
0.000, 0.167, 0.000,
0.000, 0.333, 0.000,
0.000, 0.500, 0.000,
0.000, 0.667, 0.000,
0.000, 0.833, 0.000,
0.000, 1.000, 0.000,
0.000, 0.000, 0.167,
0.000, 0.000, 0.333,
0.000, 0.000, 0.500,
0.000, 0.000, 0.667,
0.000, 0.000, 0.833,
0.000, 0.000, 1.000,
0.143, 0.143, 0.143,
0.286, 0.286, 0.286,
0.429, 0.429, 0.429,
0.571, 0.571, 0.571,
0.714, 0.714, 0.714,
0.857, 0.857, 0.857
]
).astype(np.float32)
color_list = color_list.reshape((-1, 3)) * 255
if not rgb:
color_list = color_list[:, ::-1]
return color_list
color_list = colormap()
color_list = color_list.astype('uint8').tolist()
def vis_add_mask(image, mask, color, alpha):
color = np.array(color_list[color])
mask = mask > 0.5
image[mask] = image[mask] * (1-alpha) + color * alpha
return image.astype('uint8')
def point_painter(input_image, input_points, point_color=5, point_alpha=0.9, point_radius=15, contour_color=2, contour_width=5):
h, w = input_image.shape[:2]
point_mask = np.zeros((h, w)).astype('uint8')
for point in input_points:
point_mask[point[1], point[0]] = 1
kernel = cv2.getStructuringElement(2, (point_radius, point_radius))
point_mask = cv2.dilate(point_mask, kernel)
contour_radius = (contour_width - 1) // 2
dist_transform_fore = cv2.distanceTransform(point_mask, cv2.DIST_L2, 3)
dist_transform_back = cv2.distanceTransform(1-point_mask, cv2.DIST_L2, 3)
dist_map = dist_transform_fore - dist_transform_back
# ...:::!!!:::...
contour_radius += 2
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
contour_mask = contour_mask / np.max(contour_mask)
contour_mask[contour_mask>0.5] = 1.
# paint mask
painted_image = vis_add_mask(input_image, point_mask, point_color, point_alpha)
# paint contour
painted_image = vis_add_mask(painted_image, 1-contour_mask, contour_color, 1)
return painted_image
def mask_painter(input_image, input_mask, mask_color=5, mask_alpha=0.7, contour_color=1, contour_width=3):
assert input_image.shape[:2] == input_mask.shape, 'different shape between image and mask'
# 0: background, 1: foreground
mask = np.clip(input_mask, 0, 1)
contour_radius = (contour_width - 1) // 2
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
dist_map = dist_transform_fore - dist_transform_back
# ...:::!!!:::...
contour_radius += 2
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
contour_mask = contour_mask / np.max(contour_mask)
contour_mask[contour_mask>0.5] = 1.
# paint mask
painted_image = vis_add_mask(input_image, mask, mask_color, mask_alpha)
# paint contour
painted_image = vis_add_mask(painted_image, 1-contour_mask, contour_color, 1)
return painted_image
if __name__ == '__main__':
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
input_mask = np.array(Image.open('images/painter_input_mask.jpg').convert('P'))
# example of mask painter
mask_color = 3
mask_alpha = 0.7
contour_color = 1
contour_width = 5
painted_image = mask_painter(input_image, input_mask, mask_color, mask_alpha, contour_color, contour_width)
# save
painted_image = Image.fromarray(painted_image)
painted_image.save('images/mask_painter.png')
# example of point painter
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
input_points = np.array([[500, 375], [70, 600]]) # x, y
point_color = 5
point_alpha = 0.9
point_radius = 15
contour_color = 2
contour_width = 5
painted_image_1 = point_painter(input_image, input_points, point_color, point_alpha, point_radius, contour_color, contour_width)
# save
painted_image = Image.fromarray(painted_image_1)
painted_image.save('images/point_painter_1.png')
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
painted_image_2 = point_painter(input_image, input_points, point_color=9, point_radius=20, contour_color=29)
# save
painted_image = Image.fromarray(painted_image_2)
painted_image.save('images/point_painter_2.png')

4
tracker/eval.py
View File

@@ -205,10 +205,6 @@ for vid_reader in progressbar(meta_loader, max_value=len(meta_dataset), redirect
# Run the model on this frame
prob = processor.step(rgb, msk, labels, end=(ti==vid_length-1)) # 0, background, >0, objects
# consider prob (only object channels) as prompt to refine segment results
# Upsample to original size if needed
if need_resize:
prob = F.interpolate(prob.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:,0]