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

beta -version -li

Browse Source
This commit is contained in:
memoryunreal
2023-04-19 11:34:14 +00:00
parent 579a105166
commit d7f2684303
4 changed files with 266 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files

212
app.py
View File

@@ -17,7 +17,7 @@ import torchvision
import torch
import concurrent.futures
import queue
from tools.painter import mask_painter, point_painter
# download checkpoints
def download_checkpoint(url, folder, filename):
os.makedirs(folder, exist_ok=True)
@@ -84,12 +84,18 @@ def get_frames_from_video(video_input, video_state):
"masks": [None]*len(frames),
"logits": [None]*len(frames),
"select_frame_number": 0,
"fps": 30
"fps": fps
}
return video_state, gr.update(visible=True, maximum=len(frames), value=1)
video_info = "Video Name: {}, FPS: {}, Total Frames: {}".format(video_state["video_name"], video_state["fps"], len(frames))
return video_state, video_info, gr.update(visible=True, maximum=len(frames), value=1), gr.update(visible=True, maximum=len(frames), value=1), \
gr.update(visible=True), gr.update(visible=True), \
gr.update(visible=True), gr.update(visible=True), \
gr.update(visible=True), gr.update(visible=True), \
gr.update(visible=True), gr.update(visible=True), \
gr.update(visible=True)
# get the select frame from gradio slider
def select_template(image_selection_slider, video_state):
def select_template(image_selection_slider, video_state, interactive_state):
# images = video_state[1]
image_selection_slider -= 1
@@ -100,8 +106,14 @@ def select_template(image_selection_slider, video_state):
model.samcontroler.sam_controler.reset_image()
model.samcontroler.sam_controler.set_image(video_state["origin_images"][image_selection_slider])
# # clear multi mask
# interactive_state["multi_mask"] = {"masks":[], "mask_names":[]}
return video_state["painted_images"][image_selection_slider], video_state
return video_state["painted_images"][image_selection_slider], video_state, interactive_state
def get_end_number(track_pause_number_slider, interactive_state):
interactive_state["track_end_number"] = track_pause_number_slider
return interactive_state
# use sam to get the mask
def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr.SelectData):
@@ -133,17 +145,59 @@ def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr
return painted_image, video_state, interactive_state
def add_multi_mask(video_state, interactive_state, mask_dropdown):
mask = video_state["masks"][video_state["select_frame_number"]]
interactive_state["multi_mask"]["masks"].append(mask)
interactive_state["multi_mask"]["mask_names"].append("mask_{}".format(len(interactive_state["multi_mask"]["masks"])))
return interactive_state, gr.update(choices=interactive_state["multi_mask"]["mask_names"])
def remove_multi_mask(interactive_state):
interactive_state["multi_mask"]["mask_names"]= []
interactive_state["multi_mask"]["masks"] = []
return interactive_state
def show_mask(video_state, interactive_state, mask_dropdown):
mask_dropdown.sort()
select_frame = video_state["origin_images"][video_state["select_frame_number"]]
for i in range(len(mask_dropdown)):
mask_number = int(mask_dropdown[i].split("_")[1]) - 1
mask = interactive_state["multi_mask"]["masks"][mask_number]
select_frame = mask_painter(select_frame, mask.astype('uint8'), mask_color=mask_number+2)
return select_frame
# tracking vos
def vos_tracking_video(video_state, interactive_state):
def vos_tracking_video(video_state, interactive_state, mask_dropdown):
model.xmem.clear_memory()
following_frames = video_state["origin_images"][video_state["select_frame_number"]:]
template_mask = video_state["masks"][video_state["select_frame_number"]]
if interactive_state["track_end_number"]:
following_frames = video_state["origin_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]]
else:
following_frames = video_state["origin_images"][video_state["select_frame_number"]:]
if interactive_state["multi_mask"]["masks"]:
# if mask_dropdown:
mask_dropdown.sort()
template_mask = interactive_state["multi_mask"]["masks"][int(mask_dropdown[0].split("_")[1]) - 1]
for i in range(1,len(mask_dropdown)):
mask_number = int(mask_dropdown[i].split("_")[1]) - 1
template_mask = np.clip(template_mask+interactive_state["multi_mask"]["masks"][mask_number]*(mask_number+1), 0, mask_number+1)
video_state["masks"][video_state["select_frame_number"]]= template_mask
else:
template_mask = video_state["masks"][video_state["select_frame_number"]]
fps = video_state["fps"]
masks, logits, painted_images = model.generator(images=following_frames, template_mask=template_mask)
video_state["masks"][video_state["select_frame_number"]:] = masks
video_state["logits"][video_state["select_frame_number"]:] = logits
video_state["painted_images"][video_state["select_frame_number"]:] = painted_images
if interactive_state["track_end_number"]:
video_state["masks"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = masks
video_state["logits"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = logits
video_state["painted_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = painted_images
else:
video_state["masks"][video_state["select_frame_number"]:] = masks
video_state["logits"][video_state["select_frame_number"]:] = logits
video_state["painted_images"][video_state["select_frame_number"]:] = painted_images
video_output = generate_video_from_frames(video_state["painted_images"], output_path="./result/{}".format(video_state["video_name"]), fps=fps) # import video_input to name the output video
interactive_state["inference_times"] += 1
@@ -152,7 +206,7 @@ def vos_tracking_video(video_state, interactive_state):
interactive_state["positive_click_times"]+interactive_state["negative_click_times"],
interactive_state["positive_click_times"],
interactive_state["negative_click_times"]))
#### shanggao code for mask save
if interactive_state["mask_save"]:
if not os.path.exists('./result/mask/{}'.format(video_state["video_name"].split('.')[0])):
@@ -176,6 +230,14 @@ def generate_video_from_frames(frames, output_path, fps=30):
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))
# print(output_path)
# for frame in frames:
# video.write(frame)
# video.release()
frames = torch.from_numpy(np.asarray(frames))
if not os.path.exists(os.path.dirname(output_path)):
os.makedirs(os.path.dirname(output_path))
@@ -194,7 +256,7 @@ xmem_checkpoint = download_checkpoint(xmem_checkpoint_url, folder, xmem_checkpoi
# args, defined in track_anything.py
args = parse_augment()
args.port = 12212
args.device = "cuda:4"
args.device = "cuda:1"
args.mask_save = True
model = TrackingAnything(SAM_checkpoint, xmem_checkpoint, args)
@@ -208,8 +270,15 @@ with gr.Blocks() as iface:
"inference_times": 0,
"negative_click_times" : 0,
"positive_click_times": 0,
"mask_save": args.mask_save
})
"mask_save": args.mask_save,
"multi_mask": {
"mask_names": [],
"masks": []
},
"track_end_num": None
}
)
video_state = gr.State(
{
"video_name": "",
@@ -225,43 +294,47 @@ with gr.Blocks() as iface:
with gr.Row():
# for user video input
with gr.Column(scale=1.0):
video_input = gr.Video().style(height=360)
with gr.Column():
with gr.Row(scale=0.4):
video_input = gr.Video(autosize=True)
video_info = gr.Textbox()
with gr.Row(scale=1):
# put the template frame under the radio button
with gr.Column(scale=0.5):
with gr.Column(scale=0.4):
# extract frames
with gr.Column():
extract_frames_button = gr.Button(value="Get video info", interactive=True, variant="primary")
# click points settins, negative or positive, mode continuous or single
with gr.Row():
with gr.Row(scale=0.5):
with gr.Row(scale=0.4):
point_prompt = gr.Radio(
choices=["Positive", "Negative"],
value="Positive",
label="Point Prompt",
interactive=True)
interactive=True,
visible=False)
click_mode = gr.Radio(
choices=["Continuous", "Single"],
value="Continuous",
label="Clicking Mode",
interactive=True)
interactive=True,
visible=False)
with gr.Row(scale=0.5):
clear_button_clike = gr.Button(value="Clear Clicks", interactive=True).style(height=160)
clear_button_image = gr.Button(value="Clear Image", interactive=True)
template_frame = gr.Image(type="pil",interactive=True, elem_id="template_frame").style(height=360)
image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Image Selection", invisible=False)
clear_button_click = gr.Button(value="Clear Clicks", interactive=True, visible=False).style(height=160)
Add_mask_button = gr.Button(value="Add mask", interactive=True, visible=False)
template_frame = gr.Image(type="pil",interactive=True, elem_id="template_frame", visible=False).style(height=360)
image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Image Selection", visible=False)
track_pause_number_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Track end frames", visible=False)
with gr.Column(scale=0.5):
video_output = gr.Video().style(height=360)
tracking_video_predict_button = gr.Button(value="Tracking")
with gr.Column(scale=0.4):
mask_dropdown = gr.Dropdown(multiselect=True, label="Mask_select", info=".", visible=False)
remove_mask_button = gr.Button(value="Remove mask", interactive=True, visible=False)
video_output = gr.Video(autosize=True, visible=False).style(height=360)
tracking_video_predict_button = gr.Button(value="Tracking", visible=False)
# first step: get the video information
extract_frames_button.click(
@@ -269,27 +342,51 @@ with gr.Blocks() as iface:
inputs=[
video_input, video_state
],
outputs=[video_state, image_selection_slider],
outputs=[video_state, video_info, image_selection_slider, track_pause_number_slider,point_prompt, click_mode, clear_button_click, Add_mask_button, template_frame,
tracking_video_predict_button, video_output, mask_dropdown, remove_mask_button]
)
# second step: select images from slider
image_selection_slider.release(fn=select_template,
inputs=[image_selection_slider, video_state],
outputs=[template_frame, video_state], api_name="select_image")
inputs=[image_selection_slider, video_state, interactive_state],
outputs=[template_frame, video_state, interactive_state], api_name="select_image")
track_pause_number_slider.release(fn=get_end_number,
inputs=[track_pause_number_slider, interactive_state],
outputs=[interactive_state], api_name="end_image")
# click select image to get mask using sam
template_frame.select(
fn=sam_refine,
inputs=[video_state, point_prompt, click_state, interactive_state],
outputs=[template_frame, video_state, interactive_state]
)
# add different mask
Add_mask_button.click(
fn=add_multi_mask,
inputs=[video_state, interactive_state, mask_dropdown],
outputs=[interactive_state, mask_dropdown]
)
remove_mask_button.click(
fn=remove_multi_mask,
inputs=[interactive_state],
outputs=[interactive_state]
)
# tracking video from select image and mask
tracking_video_predict_button.click(
fn=vos_tracking_video,
inputs=[video_state, interactive_state],
inputs=[video_state, interactive_state, mask_dropdown],
outputs=[video_output, video_state, interactive_state]
)
# click to get mask
mask_dropdown.change(
fn=show_mask,
inputs=[video_state, interactive_state, mask_dropdown],
outputs=[template_frame]
)
# clear input
video_input.clear(
@@ -306,10 +403,15 @@ with gr.Blocks() as iface:
"inference_times": 0,
"negative_click_times" : 0,
"positive_click_times": 0,
"mask_save": args.mask_save
"mask_save": args.mask_save,
"multi_mask": {
"mask_names": [],
"masks": []
},
"track_end_num": 0
},
[[],[]]
),
),
[],
[
video_state,
@@ -317,38 +419,10 @@ with gr.Blocks() as iface:
click_state,
],
queue=False,
show_progress=False
)
clear_button_image.click(
lambda: (
{
"origin_images": None,
"painted_images": None,
"masks": None,
"logits": None,
"select_frame_number": 0,
"fps": 30
},
{
"inference_times": 0,
"negative_click_times" : 0,
"positive_click_times": 0,
"mask_save": args.mask_save
},
[[],[]]
),
[],
[
video_state,
interactive_state,
click_state,
],
show_progress=False)
queue=False,
show_progress=False
)
clear_button_clike.click(
# points clear
clear_button_click.click(
lambda: ([[],[]]),
[],
[click_state],