def __init__(self,
parent,
ID,
img,
mask_img,
img_path,
win_xy,
color=None):
print "__init__ DrawMagicPanel"
# caculate scale of input image
self.img = img
self.mask_img = mask_img
self.img_path = img_path
height, width = self.img.shape[:2]
draw_size_w = win_xy[0] - 50
if width < draw_size_w:
self.scale = float(draw_size_w) / float(width)
draw_size_h = height * self.scale
else:
self.scale = float(draw_size_w) / float(width)
draw_size_h = height * self.scale
print self.scale
self.draw_size = (int(draw_size_w), int(draw_size_h))
wx.Panel.__init__(self, parent, ID, size=self.draw_size)
self.SetBackgroundColour("Black")
# Setup parament
self.floodmin_v = FLOOD_VALUE
self.floodmax_v = FLOOD_VALUE
maxpqueue = 30
maxnqueue = 10
# setup for floodFill
self.stack_pre = Stack()
self.stack_nex = Stack()
self.label_data = get_label_data()
self.color = tuple(self.label_data[0]['color'])
self.floodmin = (self.floodmin_v, ) * 3
self.floodmax = (self.floodmax_v, ) * 3
self.mask = np.zeros((height + 2, width + 2), np.uint8)
self.mask[:] = 0
# setup for drawbitmap
self.frame = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB) # RGB
self.bmp = wx.BitmapFromBuffer(width, height, self.frame) # RGB
# Bind for event
self.Bind(wx.EVT_PAINT, self.OnPaint)
self.Bind(wx.EVT_LEFT_DOWN, self.OnLeftDown)
self.Bind(wx.EVT_LEFT_UP, self.OnLeftUp)
self.Bind(wx.EVT_MOTION, self.OnMotion)
def __init__(self, parent, ID, img, img_path, win_xy, color=None):
# caculate scale of input image
self.img = img
self.img_path = img_path
height, width = self.img.shape[:2]
draw_size_w = win_xy[0] - 50
if width < draw_size_w:
self.scale = float(draw_size_w) / float(width)
draw_size_h = height * self.scale
else:
self.scale = float(draw_size_w) / float(width)
draw_size_h = height * self.scale
print "brush scale:", self.scale
self.draw_size = (int(draw_size_w), int(draw_size_h))
wx.Panel.__init__(self, parent, ID, size=self.draw_size)
self.SetBackgroundColour("White")
# Setup parament
self.thickness = 10
self.color = color
self.pen = wx.Pen(self.color, self.thickness, wx.SOLID)
self.lines = []
self.curLine = []
self.pos = (0, 0)
maxpqueue = 20
maxnqueue = 10
self.stack_pre = Stack()
self.stack_nex = Stack()
self.label_data = get_label_data()
# change image into wxpython bitmap
self.img = cv2.resize(self.img, self.draw_size)
self.frame = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
image = wx.EmptyImage(draw_size_w, draw_size_h)
image.SetData(self.frame.tostring())
self.bmp = image.ConvertToBitmap()
self.init_buffer()
# self.bmp = wx.BitmapFromBuffer(width, height, self.self.frame)
# Bind for event
self.Bind(wx.EVT_LEFT_DOWN, self.on_left_down)
self.Bind(wx.EVT_LEFT_UP, self.on_left_up)
self.Bind(wx.EVT_MOTION, self.on_motion)
self.Bind(wx.EVT_PAINT, self.on_paint)
def __init__(self, args: dict, split="train"):
self.args = args
self.data_root = args["data_root"]
self.mask_dir = args["mask_root"]
self.split = split
self.ref_count = args["sample_length"]
self.image_shape = self.image_width, self.image_height = (args["w"],
args["h"])
assert self.split in ["train", "val", "test"]
self.video_dict = dict()
for video_id in os.listdir(f"{self.data_root}/{split}/image"):
self.video_dict[video_id] = len(
os.listdir(f"{self.data_root}/{split}/image/{video_id}"))
self.video_ids = list(self.video_dict.keys())
self.vflipper = transforms.RandomVerticalFlip(1.)
self.hflipper = transforms.RandomHorizontalFlip(1.)
self.mask_transforms = transforms.Compose(
[Stack(), ToTorchFormatTensor()])
self.image_transforms = transforms.Compose(
[Stack(), ToTorchFormatTensor(),
Normalize(0.5, 0.5)])
def __init__(self, args: dict, split='train', debug=False):
self.args = args
self.split = split
self.sample_length = args['sample_length']
self.size = self.w, self.h = (args['w'], args['h'])
with open(os.path.join(args['data_root'], args['name'], split+'.json'), 'r') as f:
self.video_dict = json.load(f)
self.video_names = list(self.video_dict.keys())
if debug or split != 'train':
self.video_names = self.video_names[:100]
self._to_tensors = transforms.Compose([
Stack(),
ToTorchFormatTensor(), ])
def __init__(self, args: dict, split='train'):
self.args = args
self.split = split
self.sample_length = args['sample_length']
self.size = self.w, self.h = (args['w'], args['h'])
if args['name'] == 'YouTubeVOS':
vid_lst_prefix = os.path.join(args['data_root'], args['name'],
split + '_all_frames/JPEGImages')
vid_lst = os.listdir(vid_lst_prefix)
self.video_names = [
os.path.join(vid_lst_prefix, name) for name in vid_lst
]
self._to_tensors = transforms.Compose([
Stack(),
ToTorchFormatTensor(),
])
# My libs
from core.utils import Stack, ToTorchFormatTensor
parser = argparse.ArgumentParser(description="STTN")
parser.add_argument("-v", "--video", type=str, required=True)
parser.add_argument("-m", "--mask", type=str, required=True)
parser.add_argument("-c", "--ckpt", type=str, required=True)
parser.add_argument("--model", type=str, default='sttn')
args = parser.parse_args()
w, h = 432, 240
ref_length = 10
neighbor_stride = 5
default_fps = 24
_to_tensors = transforms.Compose([Stack(), ToTorchFormatTensor()])
# sample reference frames from the whole video
def get_ref_index(neighbor_ids, length):
ref_index = []
for i in range(0, length, ref_length):
if not i in neighbor_ids:
ref_index.append(i)
return ref_index
# read frame-wise masks
def read_mask(mpath):
masks = []
mnames = os.listdir(mpath)
parser.add_argument("--step", type=int, default=10)
parser.add_argument("--num_ref", type=int, default=-1)
parser.add_argument("--neighbor_stride", type=int, default=5)
parser.add_argument("--savefps", type=int, default=24)
parser.add_argument("--use_mp4", action='store_true')
args = parser.parse_args()
w, h = args.width, args.height
ref_length = args.step # ref_step
num_ref = args.num_ref
neighbor_stride = args.neighbor_stride
default_fps = args.savefps
_to_tensors = transforms.Compose([
Stack(),
ToTorchFormatTensor()])
# sample reference frames from the whole video
def get_ref_index(f, neighbor_ids, length):
ref_index = []
if num_ref == -1:
for i in range(0, length, ref_length):
if not i in neighbor_ids:
ref_index.append(i)
else:
start_idx = max(0, f - ref_length * (num_ref//2))
end_idx = min(length, f + ref_length * (num_ref//2))
for i in range(start_idx, end_idx+1, ref_length):
if not i in neighbor_ids:
class DrawMagicPanel(wx.Panel):
"""
The draw magic panel object.
# ==========================================
# == Initialisation and Window Management ==
# ==========================================
"""
def __init__(self,
parent,
ID,
img,
mask_img,
img_path,
win_xy,
color=None):
print "__init__ DrawMagicPanel"
# caculate scale of input image
self.img = img
self.mask_img = mask_img
self.img_path = img_path
height, width = self.img.shape[:2]
draw_size_w = win_xy[0] - 50
if width < draw_size_w:
self.scale = float(draw_size_w) / float(width)
draw_size_h = height * self.scale
else:
self.scale = float(draw_size_w) / float(width)
draw_size_h = height * self.scale
print self.scale
self.draw_size = (int(draw_size_w), int(draw_size_h))
wx.Panel.__init__(self, parent, ID, size=self.draw_size)
self.SetBackgroundColour("Black")
# Setup parament
self.floodmin_v = FLOOD_VALUE
self.floodmax_v = FLOOD_VALUE
maxpqueue = 30
maxnqueue = 10
# setup for floodFill
self.stack_pre = Stack()
self.stack_nex = Stack()
self.label_data = get_label_data()
self.color = tuple(self.label_data[0]['color'])
self.floodmin = (self.floodmin_v, ) * 3
self.floodmax = (self.floodmax_v, ) * 3
self.mask = np.zeros((height + 2, width + 2), np.uint8)
self.mask[:] = 0
# setup for drawbitmap
self.frame = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB) # RGB
self.bmp = wx.BitmapFromBuffer(width, height, self.frame) # RGB
# Bind for event
self.Bind(wx.EVT_PAINT, self.OnPaint)
self.Bind(wx.EVT_LEFT_DOWN, self.OnLeftDown)
self.Bind(wx.EVT_LEFT_UP, self.OnLeftUp)
self.Bind(wx.EVT_MOTION, self.OnMotion)
def OnPaint(self, event):
dc = wx.BufferedPaintDC(self)
dc.SetUserScale(self.scale, self.scale)
dc.DrawBitmap(self.bmp, 0, 0)
def OnLeftDown(self, event):
history = {'img': np.array(self.img), 'mask': np.array(self.mask)}
self.stack_pre.push(history)
self.dragline = []
self.pos = event.GetPositionTuple()
self.pos = tuple([int(x / self.scale) for x in self.pos])
self.CaptureMouse()
def OnLeftUp(self, event):
if self.HasCapture():
coords = (0, 0) + self.pos
self.dragline.append(coords)
self._floodfill(coords)
self.dragline = []
self.ReleaseMouse()
# self.stack_nex.push(history)
self.Refresh()
def OnMotion(self, event):
if event.Dragging() and event.LeftIsDown():
self.DragMotion(event)
event.Skip()
def DragMotion(self, event):
newPos = event.GetPositionTuple()
newPos = tuple([int(x / self.scale) for x in newPos])
coords = self.pos + newPos
self.dragline.append(coords)
self._floodfill(coords)
self.pos = newPos
self.Refresh()
# print "coords:", coords
def _floodfill(self, coords):
# begin floodfill
point = (coords[2], coords[3])
fill_color = (self.color[2], self.color[1], self.color[0])
# limiting mechanism 2
if coords[2] > self.draw_size[0] or coords[3] > self.draw_size[1] or \
coords[0] > self.draw_size[0] or coords[1] > self.draw_size[1]:
return
# limiting mechanism 1
if (self.img[point[1]][point[0]] == np.array(fill_color)).all():
return
# cv2.floodFill(self.mask_img, self.mask, point, fill_color,
# self.floodmin, self.floodmax)
# img = np.dstack([img, np.ones((h, w), dtype="uint8") * 255])
# cv2.addWeighted(overlay, 1.0, dst, 1.0, 0, dst)
cv2.floodFill(self.mask_img, self.mask, point, fill_color,
self.floodmin, self.floodmax)
print self.mask.shape
cv2.imwrite("test.png", self.mask_img)
# cv2.waitKey(0)
# must clean the mask otherwize cannot rechange color
# self.mask[:] = 0
self.frame = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB) # BGR
self.bmp.CopyFromBuffer(self.frame)
print "[ACTION] floodfill {} with {}".format(self.pos, self.color)
return
def set_color(self, color_rgb):
self.color = color_rgb
def get_color(self):
return self.color
def set_flood_value(self, flood_value):
self.floodmin = (flood_value, ) * 3
self.floodmax = (flood_value, ) * 3
def get_pre(self):
# Mark: canot used the self.img save into Stack() because the numpy
# is store in the memory, so we should deep copy the numpy like:
## B = np.array(A)
if not self.stack_pre.is_empty():
history = self.stack_pre.pop()
self.img = np.array(history['img'])
del self.mask
self.mask = np.array(history['mask'])
self.frame = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
self.bmp.CopyFromBuffer(self.frame)
self.Refresh()
self.Update()
def GetNext(self):
if not self.stack_nex.is_empty():
history = self.stack_nex.pop()
self.img = np.array(history['img'])
self.mask = np.array(history['mask'])
self.frame = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
self.bmp.CopyFromBuffer(self.frame)
self.Refresh()
self.Update()
def next_iamge(self, img, img_path):
del self.img
del self.frame
del self.bmp
self.img_path = img_path
self.img = np.array(img)
height, width = self.img.shape[:2]
self.mask = np.zeros((height + 2, width + 2), np.uint8)
self.frame = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB) # RGB
self.bmp = wx.BitmapFromBuffer(width, height, self.frame) # RGB
self.Refresh()
self.Update()
def save_image(self):
# self.stack_pre.clean()
# self.stack_nex.clean()
print("[Save]")
save = SaveImage(self.img, self.img_path)
self.Refresh()
self.Update()