app_save.py

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")