Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from demo import automask_image_app, automask_video_app, sahi_autoseg_app | |
| import argparse | |
| import cv2 | |
| import time | |
| from PIL import Image | |
| import numpy as np | |
| import os | |
| import sys | |
| sys.path.append(sys.path[0]+"/tracker") | |
| 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 | |
| import concurrent.futures | |
| import queue | |
| def download_checkpoint(url, folder, filename): | |
| os.makedirs(folder, exist_ok=True) | |
| filepath = os.path.join(folder, filename) | |
| if not os.path.exists(filepath): | |
| print("download checkpoints ......") | |
| response = requests.get(url, stream=True) | |
| with open(filepath, "wb") as f: | |
| for chunk in response.iter_content(chunk_size=8192): | |
| if chunk: | |
| f.write(chunk) | |
| print("download successfully!") | |
| return filepath | |
| def pause_video(play_state): | |
| print("user pause_video") | |
| play_state.append(time.time()) | |
| return play_state | |
| def play_video(play_state): | |
| print("user play_video") | |
| play_state.append(time.time()) | |
| return play_state | |
| # convert points input to prompt state | |
| 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], | |
| "input_label":click_state[1], | |
| "multimask_output":"True", | |
| } | |
| return prompt | |
| def get_frames_from_video(video_input, play_state): | |
| """ | |
| Args: | |
| video_path:str | |
| timestamp:float64 | |
| Return | |
| [[0:nearest_frame], [nearest_frame:], nearest_frame] | |
| """ | |
| video_path = video_input | |
| # video_name = video_path.split('/')[-1] | |
| try: | |
| timestamp = play_state[1] - play_state[0] | |
| except: | |
| timestamp = 0 | |
| frames = [] | |
| try: | |
| cap = cv2.VideoCapture(video_path) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if ret == True: | |
| frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) | |
| else: | |
| break | |
| except (OSError, TypeError, ValueError, KeyError, SyntaxError) as e: | |
| print("read_frame_source:{} error. {}\n".format(video_path, str(e))) | |
| # 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_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") | |
| # set image in sam when select the template frame | |
| model.samcontroler.sam_controler.set_image(nearest_frame) | |
| return frames_split, nearest_frame, nearest_frame, fps | |
| def generate_video_from_frames(frames, output_path, fps=30): | |
| """ | |
| 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 model_reset(): | |
| model.xmem.clear_memory() | |
| return None | |
| def sam_refine(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]) | |
| 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(video_state, template_mask,fps,video_input): | |
| masks, logits, painted_images = model.generator(images=video_state[1], template_mask=template_mask) | |
| video_output = generate_video_from_frames(painted_images, output_path="./output.mp4", fps=fps) | |
| # image_selection_slider = gr.Slider(minimum=1, maximum=len(video_state[1]), value=1, label="Image Selection", interactive=True) | |
| video_name = video_input.split('/')[-1].split('.')[0] | |
| result_path = os.path.join('/hhd3/gaoshang/Track-Anything/results/'+video_name) | |
| if not os.path.exists(result_path): | |
| os.makedirs(result_path) | |
| i=0 | |
| for mask in masks: | |
| np.save(os.path.join(result_path,'{:05}.npy'.format(i)), mask) | |
| i+=1 | |
| return video_output, painted_images, masks, logits | |
| def vos_tracking_image(image_selection_slider, painted_images): | |
| # images = video_state[1] | |
| percentage = image_selection_slider / 100 | |
| select_frame_num = int(percentage * len(painted_images)) | |
| return painted_images[select_frame_num], select_frame_num | |
| def interactive_correction(video_state, point_prompt, click_state, select_correction_frame, evt: gr.SelectData): | |
| """ | |
| Args: | |
| template_frame: PIL.Image | |
| point_prompt: flag for positive or negative button click | |
| click_state: [[points], [labels]] | |
| """ | |
| refine_image = video_state[1][select_correction_frame] | |
| 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) | |
| model.samcontroler.seg_again(refine_image) | |
| corrected_mask, corrected_logit, corrected_painted_image = model.first_frame_click( | |
| image=refine_image, | |
| points=np.array(prompt["input_point"]), | |
| labels=np.array(prompt["input_label"]), | |
| multimask=prompt["multimask_output"], | |
| ) | |
| return corrected_painted_image, [corrected_mask, corrected_logit, corrected_painted_image] | |
| def correct_track(video_state, select_correction_frame, corrected_state, masks, logits, painted_images, fps, video_input): | |
| model.xmem.clear_memory() | |
| # inference the following images | |
| following_images = video_state[1][select_correction_frame:] | |
| corrected_masks, corrected_logits, corrected_painted_images = model.generator(images=following_images, template_mask=corrected_state[0]) | |
| masks = masks[:select_correction_frame] + corrected_masks | |
| logits = logits[:select_correction_frame] + corrected_logits | |
| painted_images = painted_images[:select_correction_frame] + corrected_painted_images | |
| video_output = generate_video_from_frames(painted_images, output_path="./output.mp4", fps=fps) | |
| video_name = video_input.split('/')[-1].split('.')[0] | |
| result_path = os.path.join('/hhd3/gaoshang/Track-Anything/results/'+video_name) | |
| if not os.path.exists(result_path): | |
| os.makedirs(result_path) | |
| i=0 | |
| for mask in masks: | |
| np.save(os.path.join(result_path,'{:05}.npy'.format(i)), mask) | |
| i+=1 | |
| return video_output, painted_images, logits, masks | |
| # check and download checkpoints if needed | |
| SAM_checkpoint = "sam_vit_h_4b8939.pth" | |
| sam_checkpoint_url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth" | |
| xmem_checkpoint = "XMem-s012.pth" | |
| xmem_checkpoint_url = "https://github.com/hkchengrex/XMem/releases/download/v1.0/XMem-s012.pth" | |
| folder ="./checkpoints" | |
| SAM_checkpoint = download_checkpoint(sam_checkpoint_url, folder, SAM_checkpoint) | |
| xmem_checkpoint = download_checkpoint(xmem_checkpoint_url, folder, xmem_checkpoint) | |
| # args, defined in track_anything.py | |
| args = parse_augment() | |
| args.port = 12207 | |
| args.device = "cuda:5" | |
| model = TrackingAnything(SAM_checkpoint, xmem_checkpoint, args) | |
| with gr.Blocks() as iface: | |
| """ | |
| state for | |
| """ | |
| state = gr.State([]) | |
| play_state = gr.State([]) | |
| video_state = gr.State([[],[],[]]) | |
| click_state = gr.State([[],[]]) | |
| logits = gr.State([]) | |
| masks = gr.State([]) | |
| painted_images = gr.State([]) | |
| origin_image = gr.State(None) | |
| template_mask = gr.State(None) | |
| select_correction_frame = gr.State(None) | |
| corrected_state = gr.State([[],[],[]]) | |
| fps = gr.State([]) | |
| # video_name = gr.State([]) | |
| # queue value for image refresh, origin image, mask, logits, painted image | |
| with gr.Row(): | |
| # for user video input | |
| with gr.Column(scale=1.0): | |
| video_input = gr.Video().style(height=720) | |
| # listen to the user action for play and pause input video | |
| video_input.play(fn=play_video, inputs=play_state, outputs=play_state, scroll_to_output=True, show_progress=True) | |
| video_input.pause(fn=pause_video, inputs=play_state, outputs=play_state) | |
| with gr.Row(scale=1): | |
| # put the template frame under the radio button | |
| with gr.Column(scale=0.5): | |
| # click points settins, negative or positive, mode continuous or single | |
| with gr.Row(): | |
| with gr.Row(scale=0.5): | |
| point_prompt = gr.Radio( | |
| choices=["Positive", "Negative"], | |
| value="Positive", | |
| label="Point Prompt", | |
| interactive=True) | |
| click_mode = gr.Radio( | |
| choices=["Continuous", "Single"], | |
| value="Continuous", | |
| label="Clicking Mode", | |
| interactive=True) | |
| 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) | |
| with gr.Column(): | |
| template_select_button = gr.Button(value="Template select", interactive=True, variant="primary") | |
| with gr.Column(scale=0.5): | |
| # for intermedia result check and correction | |
| # 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") | |
| image_output = gr.Image(type="pil", interactive=True, elem_id="image_output").style(height=360) | |
| image_selection_slider = gr.Slider(minimum=0, maximum=100, step=0.1, value=0, label="Image Selection", interactive=True) | |
| correct_track_button = gr.Button(value="Interactive Correction") | |
| template_frame.select( | |
| fn=sam_refine, | |
| inputs=[ | |
| origin_image, point_prompt, click_state, logits | |
| ], | |
| outputs=[ | |
| template_frame, click_state, logits, template_mask | |
| ] | |
| ) | |
| template_select_button.click( | |
| fn=get_frames_from_video, | |
| inputs=[ | |
| video_input, | |
| play_state | |
| ], | |
| # outputs=[video_state, template_frame, origin_image, fps, video_name], | |
| outputs=[video_state, template_frame, origin_image, fps], | |
| ) | |
| tracking_video_predict_button.click( | |
| fn=vos_tracking_video, | |
| inputs=[video_state, template_mask, fps, video_input], | |
| outputs=[video_output, painted_images, masks, logits] | |
| ) | |
| image_selection_slider.release(fn=vos_tracking_image, | |
| inputs=[image_selection_slider, painted_images], outputs=[image_output, select_correction_frame], api_name="select_image") | |
| # correction | |
| image_output.select( | |
| fn=interactive_correction, | |
| inputs=[video_state, point_prompt, click_state, select_correction_frame], | |
| outputs=[image_output, corrected_state] | |
| ) | |
| correct_track_button.click( | |
| fn=correct_track, | |
| inputs=[video_state, select_correction_frame, corrected_state, masks, logits, painted_images, fps,video_input], | |
| outputs=[video_output, painted_images, logits, masks ] | |
| ) | |
| # clear input | |
| video_input.clear( | |
| lambda: ([], [], [[], [], []], | |
| None, "", "", "", "", "", "", "", [[],[]], | |
| None), | |
| [], | |
| [ state, play_state, video_state, | |
| template_frame, video_output, image_output, origin_image, template_mask, painted_images, masks, logits, click_state, | |
| select_correction_frame], | |
| queue=False, | |
| show_progress=False | |
| ) | |
| clear_button_image.click( | |
| fn=model_reset | |
| ) | |
| clear_button_clike.click( | |
| lambda: ([[],[]]), | |
| [], | |
| [click_state], | |
| queue=False, | |
| show_progress=False | |
| ) | |
| iface.queue(concurrency_count=1) | |
| iface.launch(debug=True, enable_queue=True, server_port=args.port, server_name="0.0.0.0") | |