Init object remover

FGT_codes

@@ -0,0 +1 @@
+Subproject commit f7bc8a2c520ef862d0a4d28a38334604eb41410c

SiamMask

@@ -0,0 +1 @@
+Subproject commit 0eaac33050fdcda81c9a25aa307fffa74c182e36

app.py

@@ -0,0 +1,259 @@
+from PIL import Image
+import gradio as gr
+from video_inpainting import video_inpainting
+from tools.test import *
+from 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 "output.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)

requirements.txt

@@ -0,0 +1,20 @@
+torch==1.10.1
+torchvision==0.11.2
+cvbase==0.5.5
+imageio==2.6.1
+matplotlib==3.1.1
+numpy==1.22.2
+opencv-python
+Pillow
+PyYAML
+scikit-image
+scipy
+tensorboardX
+imageio-ffmpeg
+Cython==0.29.34
+colorama==0.3.9
+requests==2.21.0
+fire==0.1.3
+numba==0.39.0
+h5py==2.8.0
+tqdm==4.29.1