Spaces:
Running
on
T4
Running
on
T4
| from PIL import Image | |
| import gradio as gr | |
| from FGT_codes.tool.video_inpainting import video_inpainting | |
| from SiamMask.tools.test import * | |
| from SiamMask.experiments.siammask_sharp.custom import Custom | |
| from types import SimpleNamespace | |
| import torch | |
| import numpy as np | |
| import torchvision | |
| import cv2 | |
| import sys | |
| from os.path import exists, join, basename, splitext | |
| import os | |
| project_name = './video-object-remover' | |
| sys.path.append(project_name) | |
| sys.path.append(join(project_name, 'SiamMask', | |
| 'experiments', 'siammask_sharp')) | |
| sys.path.append(join(project_name, 'SiamMask', 'models')) | |
| sys.path.append(join(project_name, 'SiamMask')) | |
| exp_path = join(project_name, 'SiamMask/experiments/siammask_sharp') | |
| pretrained_path1 = join(exp_path, 'SiamMask_DAVIS.pth') | |
| sys.path.append(join(project_name, 'FGT_codes')) | |
| sys.path.append(join(project_name, 'FGT_codes', 'tool')) | |
| sys.path.append(join(project_name, 'FGT_codes', 'LAFC', 'flowCheckPoint')) | |
| sys.path.append(join(project_name, 'FGT_codes', 'LAFC', 'checkpoint')) | |
| sys.path.append(join(project_name, 'FGT_codes', 'FGT', 'checkpoint')) | |
| sys.path.append(join(project_name, 'FGT_codes', 'LAFC', | |
| 'flowCheckPoint', 'raft-things.pth')) | |
| torch.set_grad_enabled(False) | |
| # init SiamMask | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| cfg = load_config(SimpleNamespace(config=join(exp_path, 'config_davis.json'))) | |
| siammask = Custom(anchors=cfg['anchors']) | |
| siammask = load_pretrain(siammask, pretrained_path1) | |
| siammask = siammask.eval().to(device) | |
| # constants | |
| object_x = 0 | |
| object_y = 0 | |
| object_width = 0 | |
| object_height = 0 | |
| original_frame_list = [] | |
| mask_list = [] | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--opt', default='configs/object_removal.yaml', | |
| help='Please select your config file for inference') | |
| # video completion | |
| parser.add_argument('--mode', default='object_removal', choices=[ | |
| 'object_removal', 'watermark_removal', 'video_extrapolation'], help="modes: object_removal / video_extrapolation") | |
| parser.add_argument( | |
| '--path', default='/myData/davis_resized/walking', help="dataset for evaluation") | |
| parser.add_argument( | |
| '--path_mask', default='/myData/dilateAnnotations_4/walking', help="mask for object removal") | |
| parser.add_argument( | |
| '--outroot', default='quick_start/walking3', help="output directory") | |
| parser.add_argument('--consistencyThres', dest='consistencyThres', default=5, type=float, | |
| help='flow consistency error threshold') | |
| parser.add_argument('--alpha', dest='alpha', default=0.1, type=float) | |
| parser.add_argument('--Nonlocal', dest='Nonlocal', | |
| default=False, type=bool) | |
| # RAFT | |
| parser.add_argument( | |
| '--raft_model', default='../LAFC/flowCheckPoint/raft-things.pth', help="restore checkpoint") | |
| parser.add_argument('--small', action='store_true', help='use small model') | |
| parser.add_argument('--mixed_precision', | |
| action='store_true', help='use mixed precision') | |
| parser.add_argument('--alternate_corr', action='store_true', | |
| help='use efficent correlation implementation') | |
| # LAFC | |
| parser.add_argument('--lafc_ckpts', type=str, default='../LAFC/checkpoint') | |
| # FGT | |
| parser.add_argument('--fgt_ckpts', type=str, default='../FGT/checkpoint') | |
| # extrapolation | |
| parser.add_argument('--H_scale', dest='H_scale', default=2, | |
| type=float, help='H extrapolation scale') | |
| parser.add_argument('--W_scale', dest='W_scale', default=2, | |
| type=float, help='W extrapolation scale') | |
| # Image basic information | |
| parser.add_argument('--imgH', type=int, default=256) | |
| parser.add_argument('--imgW', type=int, default=432) | |
| parser.add_argument('--flow_mask_dilates', type=int, default=8) | |
| parser.add_argument('--frame_dilates', type=int, default=0) | |
| parser.add_argument('--gpu', type=int, default=0) | |
| # FGT inference parameters | |
| 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) | |
| # visualization | |
| parser.add_argument('--vis_flows', action='store_true', | |
| help='Visualize the initialized flows') | |
| parser.add_argument('--vis_completed_flows', | |
| action='store_true', help='Visualize the completed flows') | |
| parser.add_argument('--vis_prop', action='store_true', | |
| help='Visualize the frames after stage-I filling (flow guided content propagation)') | |
| parser.add_argument('--vis_frame', action='store_true', | |
| help='Visualize frames') | |
| args = parser.parse_args() | |
| def getBoundaries(mask): | |
| if mask is None: | |
| return 0, 0, 0, 0 | |
| indexes = np.where((mask == [255, 255, 255]).all(axis=2)) | |
| print(indexes) | |
| x1 = min(indexes[1]) | |
| y1 = min(indexes[0]) | |
| x2 = max(indexes[1]) | |
| y2 = max(indexes[0]) | |
| return x1, y1, (x2-x1), (y2-y1) | |
| def track_and_mask(vid, original_frame, masked_frame): | |
| x, y, w, h = getBoundaries(masked_frame) | |
| f = 0 | |
| video_capture = cv2.VideoCapture() | |
| if video_capture.open(vid): | |
| width, height = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int( | |
| video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = video_capture.get(cv2.CAP_PROP_FPS) | |
| # can't write out mp4, so try to write into an AVI file | |
| video_writer = cv2.VideoWriter( | |
| "output.avi", cv2.VideoWriter_fourcc(*'MP42'), fps, (width, height)) | |
| video_writer2 = cv2.VideoWriter( | |
| "output_mask.avi", cv2.VideoWriter_fourcc(*'MP42'), fps, (width, height)) | |
| while video_capture.isOpened(): | |
| ret, frame = video_capture.read() | |
| if not ret: | |
| break | |
| # frame = cv2.resize(frame, (w - w % 8, h - h % 8)) | |
| if f == 0: | |
| target_pos = np.array([x + w / 2, y + h / 2]) | |
| target_sz = np.array([w, h]) | |
| # init tracker | |
| state = siamese_init( | |
| frame, target_pos, target_sz, siammask, cfg['hp'], device=device) | |
| else: | |
| # track | |
| state = siamese_track( | |
| state, frame, mask_enable=True, refine_enable=True, device=device) | |
| location = state['ploygon'].flatten() | |
| mask = state['mask'] > state['p'].seg_thr | |
| frame[:, :, 2] = (mask > 0) * 255 + \ | |
| (mask == 0) * frame[:, :, 2] | |
| mask = mask.astype(np.uint8) # convert to an unsigned byte | |
| mask = mask * 255 | |
| mask_list.append(mask) | |
| cv2.polylines(frame, [np.int0(location).reshape( | |
| (-1, 1, 2))], True, (0, 255, 0), 3) | |
| original_frame_list.append(frame) | |
| mask_list.append(mask) | |
| video_writer.write(frame) | |
| video_writer2.write(mask) | |
| f = f + 1 | |
| video_capture.release() | |
| video_writer.release() | |
| video_writer2.release() | |
| else: | |
| print("can't open the given input video file!") | |
| return "output.mp4" | |
| def inpaint_video(): | |
| video_inpainting(args, original_frame_list, mask_list) | |
| return "result.mp4" | |
| def get_first_frame(video): | |
| video_capture = cv2.VideoCapture() | |
| if video_capture.open(video): | |
| width, height = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int( | |
| video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| if video_capture.isOpened(): | |
| ret, frame = video_capture.read() | |
| RGB_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| return RGB_frame | |
| def drawRectangle(frame, mask): | |
| x1, y1, x2, y2 = getBoundaries(mask) | |
| return cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) | |
| def getStartEndPoints(mask): | |
| if mask is None: | |
| return 0, 0, 0, 0 | |
| indexes = np.where((mask == [255, 255, 255]).all(axis=2)) | |
| print(indexes) | |
| x1 = min(indexes[1]) | |
| y1 = min(indexes[0]) | |
| x2 = max(indexes[1]) | |
| y2 = max(indexes[0]) | |
| return x1, y1, x2, y2 | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| with gr.Row(): | |
| in_video = gr.Video() | |
| with gr.Row(): | |
| first_frame = gr.ImageMask() | |
| with gr.Row(): | |
| approve_mask = gr.Button(value="Approve Mask") | |
| with gr.Column(scale=1): | |
| with gr.Row(): | |
| original_image = gr.Image(interactive=False) | |
| with gr.Row(): | |
| masked_image = gr.Image(interactive=False) | |
| with gr.Column(scale=2): | |
| out_video = gr.Video() | |
| out_video_inpaint = gr.Video() | |
| track_mask = gr.Button(value="Track and Mask") | |
| inpaint = gr.Button(value="Inpaint") | |
| in_video.change(fn=get_first_frame, inputs=[ | |
| in_video], outputs=[first_frame]) | |
| approve_mask.click(lambda x: [x['image'], x['mask']], first_frame, [ | |
| original_image, masked_image]) | |
| track_mask.click(fn=track_and_mask, inputs=[ | |
| in_video, original_image, masked_image], outputs=[out_video]) | |
| inpaint.click(fn=inpaint_video, outputs=[out_video_inpaint]) | |
| demo.launch(share=True, debug=True) | |