|
|
import cv2
|
|
|
import numpy as np
|
|
|
import psutil
|
|
|
import torch
|
|
|
from XMem.inference.inference_core import InferenceCore
|
|
|
from XMem.model.network import XMem
|
|
|
from XMem.inference.data.mask_mapper import MaskMapper
|
|
|
from config import XMEM_CONFIG, DEVICE
|
|
|
from torchvision import transforms
|
|
|
from XMem.util.range_transform import im_normalization
|
|
|
from XMem.inference.interact.interactive_utils import overlay_davis
|
|
|
from segmenter import Segmenter
|
|
|
from tools.mask_display import visualize_unique_mask, visualize_wb_mask, mask_map
|
|
|
from tools.contour_detector import getting_coordinates
|
|
|
|
|
|
|
|
|
class TrackerCore:
|
|
|
|
|
|
name_version = 'XMem'
|
|
|
|
|
|
def __init__(self, device: str = DEVICE):
|
|
|
self.device = device
|
|
|
if self.device.lower() != 'cpu':
|
|
|
self.network = XMem(XMEM_CONFIG, 'checkpoints/XMem.pth').eval().to('cuda')
|
|
|
else:
|
|
|
self.network = XMem(
|
|
|
XMEM_CONFIG, 'checkpoints/XMem.pth', map_location='cpu'
|
|
|
).eval()
|
|
|
self.processor = InferenceCore(self.network, XMEM_CONFIG)
|
|
|
|
|
|
self.im_transform = transforms.Compose(
|
|
|
[transforms.ToTensor(), im_normalization]
|
|
|
)
|
|
|
self.mapper = MaskMapper()
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def track(
|
|
|
self, frame: np.ndarray, mask_segmet: np.ndarray = None, exhaustive=False
|
|
|
):
|
|
|
if mask_segmet is not None:
|
|
|
mask, labels = self.mapper.convert_mask(mask_segmet, exhaustive)
|
|
|
mask = torch.Tensor(mask).to(self.device)
|
|
|
self.processor.set_all_labels(list(self.mapper.remappings.values()))
|
|
|
else:
|
|
|
mask = None
|
|
|
labels = None
|
|
|
|
|
|
frame_tensor = self.im_transform(frame).to(self.device)
|
|
|
probs = self.processor.step(frame_tensor, mask, labels)
|
|
|
|
|
|
out_mask = torch.argmax(probs, dim=0)
|
|
|
out_mask = (out_mask.detach().cpu().numpy()).astype(np.uint8)
|
|
|
final_mask = np.zeros_like(out_mask)
|
|
|
|
|
|
|
|
|
for k, v in self.mapper.remappings.items():
|
|
|
final_mask[out_mask == v] = k
|
|
|
|
|
|
return final_mask
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def clear_memory(self):
|
|
|
self.processor.clear_memory()
|
|
|
self.mapper.clear_lables()
|
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
path = 'video-test/video.mp4'
|
|
|
video = cv2.VideoCapture(path)
|
|
|
ret, frame = video.read()
|
|
|
frame_cop = frame.copy()
|
|
|
video.release()
|
|
|
bboxes = [(476, 166, 102, 154), (8, 252, 91, 149), (106, 335, 211, 90)]
|
|
|
points = [[531, 230], [45, 321], [226, 360]]
|
|
|
seg = Segmenter('models/FastSAM-x.pt')
|
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
seg.prompt = frame
|
|
|
seg.get_mask_by_box_prompt(bboxes)
|
|
|
mask, unique_mask = seg.convert_mask_to_color()
|
|
|
seg.clear_memory()
|
|
|
masks = []
|
|
|
images = []
|
|
|
traker = TrackerCore()
|
|
|
frames_to_propagate = 200
|
|
|
current_frame_index = 0
|
|
|
cap = cv2.VideoCapture(path)
|
|
|
while cap.isOpened():
|
|
|
current_memory_usage = psutil.virtual_memory().percent
|
|
|
|
|
|
if current_memory_usage > 90:
|
|
|
break
|
|
|
ret, frame_v = cap.read()
|
|
|
if not ret:
|
|
|
break
|
|
|
if current_frame_index > frames_to_propagate:
|
|
|
break
|
|
|
|
|
|
if current_frame_index == 0:
|
|
|
mask = traker.track(frame_v, unique_mask)
|
|
|
masks.append(mask)
|
|
|
images.append(frame_v)
|
|
|
else:
|
|
|
mask = traker.track(frame_v)
|
|
|
masks.append(mask)
|
|
|
images.append(frame_v)
|
|
|
|
|
|
current_frame_index += 1
|
|
|
video.release()
|
|
|
im3 = visualize_wb_mask(masks[200])
|
|
|
ima = images[200].copy()
|
|
|
for m in mask_map(masks[200]):
|
|
|
for box in getting_coordinates(m):
|
|
|
(x, y, w, h) = [v for v in box]
|
|
|
cv2.rectangle(ima, (x, y), (x + w, y + h), (0, 255, 0), 2)
|
|
|
image_m = overlay_davis(images[200], masks[200])
|
|
|
cv2.imshow('image_m200', image_m)
|
|
|
cv2.imshow('ima_rect', ima)
|
|
|
cv2.imshow('im2', im3)
|
|
|
cv2.waitKey(0)
|
|
|
cv2.destroyAllWindows()
|
|
|
|
