Spaces:
Runtime error
Runtime error
File size: 7,264 Bytes
5fd3fc6 | 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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | from collections import deque
import numpy as np
import torch
from lib.MiVOS_STCN.inference_core import InferenceCore
from lib.MiVOS_STCN.interact.interactive_utils import images_to_torch, load_images
from lib.MiVOS_STCN.interact.s2m_controller import S2MController
from lib.MiVOS_STCN.model.fusion_net import FusionNet
from lib.MiVOS_STCN.model.propagation.prop_net import PropagationNetwork
from lib.MiVOS_STCN.model.s2m.s2m_network import deeplabv3plus_resnet50 as S2M
from util.scribble_util import MyScribbleInteraction
class MiVOS_Manager:
def __init__(self, image_folder, num_objects=1):
# Check if CUDA is available
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Set up models
prop_saved = torch.load('saves/stcn.pth', map_location=device)
prop_model = PropagationNetwork().to(device).eval()
prop_model.load_state_dict(prop_saved)
fusion_saved = torch.load('saves/fusion_stcn.pth', map_location=device)
fuse_model = FusionNet().to(device).eval()
fuse_model.load_state_dict(fusion_saved)
s2m_saved = torch.load('saves/s2m.pth', map_location=device)
s2m_model = S2M().to(device).eval()
s2m_model.load_state_dict(s2m_saved)
# Loads the images/masks
# Set resolution=-1 to use original size
self.images = load_images(image_folder)
self.num_frames, self.height, self.width = self.images.shape[:3]
self.num_objects = num_objects
self.s2m_controller = S2MController(s2m_model, num_objects=self.num_objects, ignore_class=255, device=device)
self.processor = InferenceCore(prop_model, fuse_model, images_to_torch(self.images, device=device),
self.num_objects, mem_freq=5, mem_profile=0, device=device)
# initialize visualization
self.current_mask = np.zeros((self.num_frames, self.height, self.width), dtype=np.uint8)
self.vis_map = np.zeros((self.height, self.width, 3), dtype=np.uint8)
self.vis_alpha = np.zeros((self.height, self.width, 1), dtype=np.float32)
self.brush_vis_map = np.zeros((self.height, self.width, 3), dtype=np.uint8)
self.brush_vis_alpha = np.zeros((self.height, self.width, 1), dtype=np.float32)
self.vis_hist = deque(maxlen=100)
# self.cursur is set to the number of the current frame we look at
self.cursur = 0
# initialize action
self.interactions = {}
self.interactions['interact'] = [[] for _ in range(self.num_frames)]
self.interactions['annotated_frame'] = []
self.this_frame_interactions = []
self.interaction = None
self.reset_this_interaction()
self.interacted_mask = None
def clear_visualization(self):
self.vis_map.fill(0)
self.vis_alpha.fill(0)
self.vis_hist.clear()
self.vis_hist.append((self.vis_map.copy(), self.vis_alpha.copy()))
def reset_this_interaction(self):
self.complete_interaction()
self.clear_visualization()
self.interaction = None
self.this_frame_interactions = []
def on_run(self):
"""
Propagate the masks
"""
if self.interacted_mask is None:
print('Cannot propagate! No interacted mask!')
return
# Create a list of propagated masks
with torch.cuda.amp.autocast(enabled=True) and torch.set_grad_enabled(False):
self.current_mask = self.processor.interact(self.interacted_mask, self.cursur)
self.interacted_mask = None
self.reset_this_interaction()
print('Propagation finished')
return self.current_mask
def on_drawn(self, drawing_points, frame_num, k):
"""
Execute after a scribble was drawn
:param drawing_points: The points in the scribble as a list of tuples (only points, no connection lines)
:param frame_num: The number of the frame that is being drawn on
:param k: Object id, 0 for negative scribbles, else 1
:return: Predicted mask of the scribbles
"""
self.cursur = frame_num
# Push last vis map into history
self.vis_hist.append((self.vis_map.copy(), self.vis_alpha.copy()))
prev_hard_mask = self.processor.masks[self.cursur]
image = self.processor.images[:, self.cursur]
h, w = self.height, self.width
if self.interaction is None:
self.interaction = MyScribbleInteraction(image, prev_hard_mask, (h, w), self.s2m_controller,
self.num_objects)
self.interaction.push_drawing(drawing_points, k)
interaction = self.interaction
interaction.end_path()
with torch.cuda.amp.autocast(enabled=True) and torch.set_grad_enabled(False):
self.interacted_mask = interaction.predict()
return self.update_interacted_mask()
def on_undo(self):
if self.interaction is None:
if len(self.this_frame_interactions) > 1:
self.this_frame_interactions = self.this_frame_interactions[:-1]
self.interacted_mask = self.this_frame_interactions[-1].predict()
else:
self.reset_this_interaction()
self.interacted_mask = self.processor.prob[:, self.cursur].clone()
else:
if self.interaction.can_undo():
self.interacted_mask = self.interaction.undo()
else:
if len(self.this_frame_interactions) > 0:
self.interaction = None
self.interacted_mask = self.this_frame_interactions[-1].predict()
else:
self.reset_this_interaction()
self.interacted_mask = self.processor.prob[:, self.cursur].clone()
# Update visualization
if len(self.vis_hist) > 0:
# Might be empty if we are undoing the entire interaction
self.vis_map, self.vis_alpha = self.vis_hist.pop()
# Commit changes
return self.update_interacted_mask()
def on_reset(self):
# DO not edit prob -- we still need the mask diff
self.processor.masks[self.cursur].zero_()
self.processor.np_masks[self.cursur].fill(0)
self.current_mask[self.cursur].fill(0)
self.reset_this_interaction()
# return self.current_mask[self.cursur]
def update_interacted_mask(self):
"""
Calculate the currently interacted mask
:return: a transparent image with the mask in red
"""
self.processor.update_mask_only(self.interacted_mask, self.cursur)
self.current_mask[self.cursur] = self.processor.np_masks[self.cursur]
return self.current_mask[self.cursur]
def complete_interaction(self):
if self.interaction is not None:
self.clear_visualization()
self.interactions['annotated_frame'].append(self.cursur)
self.interactions['interact'][self.cursur].append(self.interaction)
self.this_frame_interactions.append(self.interaction)
self.interaction = None
def get_size(self):
return self.height, self.width
|