mirror of
https://github.com/gaomingqi/Track-Anything.git
synced 2025-12-16 08:27:49 +01:00
Merge branch 'master' of https://github.com/gaomingqi/VOS-Anything
This commit is contained in:
gaomingqi
128
app.py
128
app.py
@@ -12,6 +12,9 @@ sys.path.append(sys.path[0]+"/tracker/model")
|
||||
from track_anything import TrackingAnything
|
||||
from track_anything import parse_augment
|
||||
import requests
|
||||
import json
|
||||
import torchvision
|
||||
import torch
|
||||
|
||||
def download_checkpoint(url, folder, filename):
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
@@ -41,8 +44,7 @@ xmem_checkpoint = download_checkpoint(xmem_checkpoint_url, folder, xmem_checkpoi
|
||||
|
||||
# args, defined in track_anything.py
|
||||
args = parse_augment()
|
||||
args.port=12212
|
||||
|
||||
args.port = 12213
|
||||
model = TrackingAnything(SAM_checkpoint, xmem_checkpoint, args)
|
||||
|
||||
|
||||
@@ -60,13 +62,15 @@ def play_video(play_state):
|
||||
return play_state
|
||||
|
||||
# convert points input to prompt state
|
||||
def get_prompt(inputs, click_state):
|
||||
points = []
|
||||
labels = []
|
||||
def get_prompt(click_state, click_input):
|
||||
inputs = json.loads(click_input)
|
||||
points = click_state[0]
|
||||
labels = click_state[1]
|
||||
for input in inputs:
|
||||
points.append(input[:2])
|
||||
labels.append(input[2])
|
||||
click_state[0] = points
|
||||
click_state[1] = labels
|
||||
prompt = {
|
||||
"prompt_type":["click"],
|
||||
"input_point":click_state[0],
|
||||
@@ -99,22 +103,95 @@ def get_frames_from_video(video_input, play_state):
|
||||
except (OSError, TypeError, ValueError, KeyError, SyntaxError) as e:
|
||||
print("read_frame_source:{} error. {}\n".format(video_path, str(e)))
|
||||
|
||||
for frame in frames:
|
||||
frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
for index, frame in enumerate(frames):
|
||||
frames[index] = np.asarray(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
|
||||
|
||||
key_frame_index = int(timestamp * fps)
|
||||
nearest_frame = frames[key_frame_index]
|
||||
frames = [frames[:key_frame_index], frames[key_frame_index:], nearest_frame]
|
||||
return frames, nearest_frame
|
||||
frames_split = [frames[:key_frame_index], frames[key_frame_index:], nearest_frame]
|
||||
# output_path='./seperate.mp4'
|
||||
# torchvision.io.write_video(output_path, frames[1], fps=fps, video_codec="libx264")
|
||||
|
||||
def inference_all(template_frame, evt:gr.SelectData):
|
||||
coordinate = "[[{},{},1]]".format(evt.index[0], evt.index[1])
|
||||
# set image in sam when select the template frame
|
||||
model.samcontroler.sam_controler.set_image(nearest_frame)
|
||||
return frames_split, nearest_frame, nearest_frame
|
||||
|
||||
def generate_video_from_frames(frames, output_path, fps=10):
|
||||
"""
|
||||
Generates a video from a list of frames.
|
||||
|
||||
Args:
|
||||
frames (list of numpy arrays): The frames to include in the video.
|
||||
output_path (str): The path to save the generated video.
|
||||
fps (int, optional): The frame rate of the output video. Defaults to 30.
|
||||
"""
|
||||
# height, width, layers = frames[0].shape
|
||||
# fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
# video = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
|
||||
|
||||
# for frame in frames:
|
||||
# video.write(frame)
|
||||
|
||||
# video.release()
|
||||
frames = torch.from_numpy(np.asarray(frames))
|
||||
output_path='./output.mp4'
|
||||
torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
|
||||
return output_path
|
||||
|
||||
# def get_video_from_frames():
|
||||
|
||||
# return video_output
|
||||
|
||||
def inference_all(origin_frame, point_prompt, click_state, logit, evt:gr.SelectData):
|
||||
"""
|
||||
Args:
|
||||
template_frame: PIL.Image
|
||||
point_prompt: flag for positive or negative button click
|
||||
click_state: [[points], [labels]]
|
||||
"""
|
||||
if point_prompt == "Positive":
|
||||
coordinate = "[[{},{},1]]".format(evt.index[0], evt.index[1])
|
||||
else:
|
||||
coordinate = "[[{},{},0]]".format(evt.index[0], evt.index[1])
|
||||
|
||||
# prompt for sam model
|
||||
prompt = get_prompt(click_state=click_state, click_input=coordinate)
|
||||
|
||||
# default value
|
||||
points = np.array([[evt.index[0],evt.index[1]]])
|
||||
labels= np.array([1])
|
||||
mask, logit, painted_image = model.inference_step(first_flag=True, interact_flag=False, image=np.asarray(template_frame), same_image_flag=False,points=points, labels=labels,logits=None,multimask=True)
|
||||
return painted_image
|
||||
# points = np.array([[evt.index[0],evt.index[1]]])
|
||||
# labels= np.array([1])
|
||||
if len(logit)==0:
|
||||
logit = None
|
||||
|
||||
mask, logit, painted_image = model.first_frame_click(
|
||||
image=origin_frame,
|
||||
points=np.array(prompt["input_point"]),
|
||||
labels=np.array(prompt["input_label"]),
|
||||
multimask=prompt["multimask_output"],
|
||||
|
||||
)
|
||||
return painted_image, click_state, logit, mask
|
||||
|
||||
def vos_tracking(video_state, template_mask):
|
||||
|
||||
masks, logits, painted_images = model.generator(images=video_state[1], mask=template_mask)
|
||||
video_output = generate_video_from_frames(painted_images, output_path="./output.mp4")
|
||||
return video_output
|
||||
|
||||
# upload file
|
||||
# def upload_callback(image_input, state):
|
||||
# state = [] + [('Image size: ' + str(image_input.size), None)]
|
||||
# click_state = [[], [], []]
|
||||
# res = 1024
|
||||
# width, height = image_input.size
|
||||
# ratio = min(1.0 * res / max(width, height), 1.0)
|
||||
# if ratio < 1.0:
|
||||
# image_input = image_input.resize((int(width * ratio), int(height * ratio)))
|
||||
# print('Scaling input image to {}'.format(image_input.size))
|
||||
# model.segmenter.image = None
|
||||
# model.segmenter.image_embedding = None
|
||||
# model.segmenter.set_image(image_input)
|
||||
# return state, state, image_input, click_state, image_input
|
||||
|
||||
|
||||
with gr.Blocks() as iface:
|
||||
@@ -125,6 +202,9 @@ with gr.Blocks() as iface:
|
||||
play_state = gr.State([])
|
||||
video_state = gr.State([[],[],[]])
|
||||
click_state = gr.State([[],[]])
|
||||
logits = gr.State([])
|
||||
origin_image = gr.State(None)
|
||||
template_mask = gr.State(None)
|
||||
|
||||
with gr.Row():
|
||||
|
||||
@@ -166,7 +246,9 @@ with gr.Blocks() as iface:
|
||||
|
||||
|
||||
# for intermedia result check and correction
|
||||
intermedia_image = gr.Image(type="pil", interactive=True, elem_id="intermedia_frame").style(height=360)
|
||||
# intermedia_image = gr.Image(type="pil", interactive=True, elem_id="intermedia_frame").style(height=360)
|
||||
video_output = gr.Video().style(height=360)
|
||||
tracking_video_predict_button = gr.Button(value="Tracking")
|
||||
|
||||
# seg_automask_video_points_per_batch = gr.Slider(
|
||||
# minimum=0,
|
||||
@@ -176,7 +258,7 @@ with gr.Blocks() as iface:
|
||||
# label="Points per Batch",
|
||||
# )
|
||||
|
||||
seg_automask_video_predict = gr.Button(value="Generator")
|
||||
|
||||
|
||||
|
||||
# Display the first frame
|
||||
@@ -207,20 +289,24 @@ with gr.Blocks() as iface:
|
||||
video_input,
|
||||
play_state
|
||||
],
|
||||
outputs=[video_state, template_frame],
|
||||
outputs=[video_state, template_frame, origin_image],
|
||||
)
|
||||
|
||||
template_frame.select(
|
||||
fn=inference_all,
|
||||
inputs=[
|
||||
template_frame
|
||||
origin_image, point_prompt, click_state, logits
|
||||
],
|
||||
outputs=[
|
||||
template_frame
|
||||
template_frame, click_state, logits, template_mask
|
||||
]
|
||||
|
||||
)
|
||||
|
||||
tracking_video_predict_button.click(
|
||||
fn=vos_tracking,
|
||||
inputs=[video_state, template_mask],
|
||||
outputs=[video_output]
|
||||
)
|
||||
# clear
|
||||
# clear_button_clike.click(
|
||||
# lambda x: ([[], [], []], x, ""),
|
||||
|
||||
@@ -31,6 +31,7 @@ class BaseSegmenter:
|
||||
def set_image(self, image: np.ndarray):
|
||||
# PIL.open(image_path) 3channel: RGB
|
||||
# image embedding: avoid encode the same image multiple times
|
||||
self.orignal_image = image
|
||||
if self.embedded:
|
||||
print('repeat embedding, please reset_image.')
|
||||
return
|
||||
|
||||
@@ -45,18 +45,39 @@ class SamControler():
|
||||
self.sam_controler.set_image(image)
|
||||
return
|
||||
|
||||
|
||||
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True):
|
||||
'''
|
||||
it is used in first frame in video
|
||||
return: mask, logit, painted image(mask+point)
|
||||
'''
|
||||
self.sam_controler.set_image(image)
|
||||
prompts = {
|
||||
'point_coords': points,
|
||||
'point_labels': labels,
|
||||
}
|
||||
masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask)
|
||||
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
|
||||
# self.sam_controler.set_image(image)
|
||||
origal_image = self.sam_controler.orignal_image
|
||||
neg_flag = labels[-1]
|
||||
if neg_flag==1:
|
||||
#find neg
|
||||
prompts = {
|
||||
'point_coords': points,
|
||||
'point_labels': labels,
|
||||
}
|
||||
masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask)
|
||||
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
|
||||
prompts = {
|
||||
'point_coords': points,
|
||||
'point_labels': labels,
|
||||
'mask_input': logit[None, :, :]
|
||||
}
|
||||
masks, scores, logits = self.sam_controler.predict(prompts, 'both', multimask)
|
||||
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
|
||||
else:
|
||||
#find positive
|
||||
prompts = {
|
||||
'point_coords': points,
|
||||
'point_labels': labels,
|
||||
}
|
||||
masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask)
|
||||
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
|
||||
|
||||
|
||||
assert len(points)==len(labels)
|
||||
|
||||
@@ -68,6 +89,7 @@ class SamControler():
|
||||
return mask, logit, painted_image
|
||||
|
||||
def interact_loop(self, image:np.ndarray, same: bool, points:np.ndarray, labels: np.ndarray, logits: np.ndarray=None, multimask=True):
|
||||
origal_image = self.sam_controler.orignal_image
|
||||
if same:
|
||||
'''
|
||||
true; loop in the same image
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
import sys
|
||||
sys.path.append("/hhd3/gaoshang/Track-Anything/tracker")
|
||||
import PIL
|
||||
from tools.interact_tools import SamControler
|
||||
from tracker.base_tracker import BaseTracker
|
||||
import numpy as np
|
||||
@@ -24,7 +27,34 @@ class TrackingAnything():
|
||||
|
||||
mask, logit, painted_image = self.xmem.track(image, logit)
|
||||
return mask, logit, painted_image
|
||||
|
||||
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True):
|
||||
mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask)
|
||||
return mask, logit, painted_image
|
||||
|
||||
def interact(self, image: np.ndarray, same_image_flag: bool, points:np.ndarray, labels: np.ndarray, logits: np.ndarray=None, multimask=True):
|
||||
mask, logit, painted_image = self.samcontroler.interact_loop(image, same_image_flag, points, labels, logits, multimask)
|
||||
return mask, logit, painted_image
|
||||
|
||||
def generator(self, images: list, mask:np.ndarray):
|
||||
|
||||
masks = []
|
||||
logits = []
|
||||
painted_images = []
|
||||
for i in range(len(images)):
|
||||
if i ==0:
|
||||
mask, logit, painted_image = self.xmem.track(images[i], mask)
|
||||
masks.append(mask)
|
||||
logits.append(logit)
|
||||
painted_images.append(painted_image)
|
||||
|
||||
else:
|
||||
mask, logit, painted_image = self.xmem.track(images[i])
|
||||
masks.append(mask)
|
||||
logits.append(logit)
|
||||
painted_images.append(painted_image)
|
||||
return masks, logits, painted_images
|
||||
|
||||
|
||||
def parse_augment():
|
||||
parser = argparse.ArgumentParser()
|
||||
@@ -36,4 +66,27 @@ def parse_augment():
|
||||
|
||||
if args.debug:
|
||||
print(args)
|
||||
return args
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
masks = None
|
||||
logits = None
|
||||
painted_images = None
|
||||
images = []
|
||||
image = np.array(PIL.Image.open('/hhd3/gaoshang/truck.jpg'))
|
||||
args = parse_augment()
|
||||
# images.append(np.ones((20,20,3)).astype('uint8'))
|
||||
# images.append(np.ones((20,20,3)).astype('uint8'))
|
||||
images.append(image)
|
||||
images.append(image)
|
||||
|
||||
mask = np.zeros_like(image)[:,:,0]
|
||||
mask[0,0]= 1
|
||||
trackany = TrackingAnything('/ssd1/gaomingqi/checkpoints/sam_vit_h_4b8939.pth','/ssd1/gaomingqi/checkpoints/XMem-s012.pth', args)
|
||||
masks, logits ,painted_images= trackany.generator(images, mask)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user