Spaces:
Runtime error
Runtime error
File size: 5,032 Bytes
7f86517 acc3d84 7f86517 af32f74 7f86517 945ea31 90db91d 7f86517 e1eb47f 7f86517 59e5beb 0ca8766 7f86517 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
import atexit
import bisect
import multiprocessing as mp
from collections import deque
import cv2
import torch
import argparse
import glob
import multiprocessing as mp
import numpy as np
import os
import tempfile
import time
import warnings
import cv2
import subprocess
import tqdm
import gradio as gr
TOTAL_FRAMES = 60
subprocess.run(["pip", "install", "git+https://github.com/wjf5203/VNext.git"])
subprocess.run(["git", "clone", "https://github.com/wjf5203/VNext"])
from detectron2.data import MetadataCatalog
from detectron2.data.detection_utils import read_image
from detectron2.engine.defaults import DefaultPredictor
from detectron2.utils.video_visualizer import VideoVisualizer
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.config import get_cfg
from detectron2.utils.logger import setup_logger
def test_opencv_video_format(codec, file_ext):
with tempfile.TemporaryDirectory(prefix="video_format_test") as dir:
filename = os.path.join(dir, "test_file" + file_ext)
writer = cv2.VideoWriter(
filename=filename,
fourcc=cv2.VideoWriter_fourcc(*codec),
fps=float(30),
frameSize=(10, 10),
isColor=True,
)
[writer.write(np.zeros((10, 10, 3), np.uint8)) for _ in range(30)]
writer.release()
if os.path.isfile(filename):
return True
return False
def setup_cfg(cfg):
# load config from file and command-line arguments
cfg = get_cfg()
# To use demo for Panoptic-DeepLab, please uncomment the following two lines.
# from detectron2.projects.panoptic_deeplab import add_panoptic_deeplab_config # noqa
# add_panoptic_deeplab_config(cfg)
cfg.merge_from_file("VNext/configs/quick_schedules/mask_rcnn_R_50_FPN_inference_acc_test.yaml")
# Set score_threshold for builtin models
cfg.MODEL.RETINANET.SCORE_THRESH_TEST = 0.5
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = 0.5
cfg.freeze()
return cfg
predictor = DefaultPredictor(setup_cfg({}))
metadata = MetadataCatalog.get("__unused")
def run_on_video(video, total_frames):
video_visualizer = VideoVisualizer(metadata, ColorMode.IMAGE)
def _frame_from_video(video):
while video.isOpened():
success, frame = video.read()
if success:
yield frame
else:
break
def process_predictions(frame, predictions):
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
if "panoptic_seg" in predictions:
panoptic_seg, segments_info = predictions["panoptic_seg"]
vis_frame = video_visualizer.draw_panoptic_seg_predictions(
frame, panoptic_seg.to("cpu"), segments_info
)
elif "instances" in predictions:
predictions = predictions["instances"].to("cpu")
vis_frame = video_visualizer.draw_instance_predictions(frame, predictions)
elif "sem_seg" in predictions:
vis_frame = video_visualizer.draw_sem_seg(
frame, predictions["sem_seg"].argmax(dim=0).to("cpu")
)
# Converts Matplotlib RGB format to OpenCV BGR format
vis_frame = cv2.cvtColor(vis_frame.get_image(), cv2.COLOR_RGB2BGR)
return vis_frame
frame_gen = _frame_from_video(video)
i = 0
for frame in frame_gen:
i += 1
if i == total_frames:
return
yield process_predictions(frame, predictor(frame))
def inference(video):
video = cv2.VideoCapture(video)
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
frames_per_second = video.get(cv2.CAP_PROP_FPS)
num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
print(num_frames)
if num_frames>TOTAL_FRAMES:
num_frames=TOTAL_FRAMES
codec, file_ext = (
("x264", ".mkv") if test_opencv_video_format("x264", ".mkv") else ("mp4v", ".mp4")
)
print(codec, file_ext)
output_fname = "result.mp4"
output_file = cv2.VideoWriter(
filename=output_fname,
fourcc=cv2.VideoWriter_fourcc(*codec),
fps=float(frames_per_second),
frameSize=(width, height),
isColor=True,
)
for vis_frame in tqdm.tqdm(run_on_video(video, num_frames), total=num_frames):
output_file.write(vis_frame)
video.release()
output_file.release()
out_file = tempfile.NamedTemporaryFile(suffix="out.mp4", delete=False)
subprocess.run(f"ffmpeg -y -loglevel quiet -stats -i {output_fname} -c:v libx264 {out_file.name}".split())
return out_file.name
video_interface = gr.Interface(
fn=inference,
inputs=[
gr.Video(type="file"),
],
outputs=gr.Video(type="file", format="mp4"),
examples=[
["inps.mp4"], ["example_3.mp4"],
],
allow_flagging=False,
allow_screenshot=False,
title="VNext",
description="demo for <a href='https://github.com/wjf5203/VNext'>wjf5203/VNext</a>"
).launch(debug=True)
|