Spaces:
Runtime error
Runtime error
isLinXu
commited on
Commit
·
0683257
1
Parent(s):
bb59176
update
Browse files- app.py +802 -0
- requirements.txt +19 -0
app.py
ADDED
|
@@ -0,0 +1,802 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
os.system("git clone https://github.com/IDEA-Research/detrex.git")
|
| 3 |
+
# os.system("python -m pip install git+https://github.com/facebookresearch/detectron2.git@v0.6#egg=detectron2")
|
| 4 |
+
os.system("python -m pip install git+https://github.com/IDEA-Research/detrex.git@v0.5.0#egg=detrex")
|
| 5 |
+
os.system("git submodule sync")
|
| 6 |
+
os.system("git submodule update --init")
|
| 7 |
+
# os.system("cd detrex && pip install -e .")
|
| 8 |
+
# os.system("pip install fairscale")
|
| 9 |
+
|
| 10 |
+
import argparse
|
| 11 |
+
import glob
|
| 12 |
+
import multiprocessing as mp
|
| 13 |
+
import numpy as np
|
| 14 |
+
import os
|
| 15 |
+
import sys
|
| 16 |
+
import tempfile
|
| 17 |
+
import time
|
| 18 |
+
import warnings
|
| 19 |
+
import cv2
|
| 20 |
+
import torch
|
| 21 |
+
import tqdm
|
| 22 |
+
import gradio as gr
|
| 23 |
+
|
| 24 |
+
# from demo.predictors import VisualizationDemo
|
| 25 |
+
from detectron2.checkpoint import DetectionCheckpointer
|
| 26 |
+
from detectron2.config import LazyConfig, instantiate
|
| 27 |
+
from detectron2.data.detection_utils import read_image
|
| 28 |
+
from detectron2.utils.logger import setup_logger
|
| 29 |
+
|
| 30 |
+
import warnings
|
| 31 |
+
|
| 32 |
+
warnings.filterwarnings("ignore")
|
| 33 |
+
|
| 34 |
+
import atexit
|
| 35 |
+
import bisect
|
| 36 |
+
from copy import copy
|
| 37 |
+
import multiprocessing as mp
|
| 38 |
+
from collections import deque
|
| 39 |
+
import cv2
|
| 40 |
+
import torch
|
| 41 |
+
|
| 42 |
+
import detectron2.data.transforms as T
|
| 43 |
+
from detectron2.data import MetadataCatalog
|
| 44 |
+
from detectron2.structures import Instances
|
| 45 |
+
from detectron2.utils.video_visualizer import VideoVisualizer
|
| 46 |
+
from detectron2.utils.visualizer import ColorMode, Visualizer
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def filter_predictions_with_confidence(predictions, confidence_threshold=0.5):
|
| 50 |
+
if "instances" in predictions:
|
| 51 |
+
preds = predictions["instances"]
|
| 52 |
+
keep_idxs = preds.scores > confidence_threshold
|
| 53 |
+
predictions = copy(predictions) # don't modify the original
|
| 54 |
+
predictions["instances"] = preds[keep_idxs]
|
| 55 |
+
return predictions
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
class VisualizationDemo(object):
|
| 59 |
+
def __init__(
|
| 60 |
+
self,
|
| 61 |
+
model,
|
| 62 |
+
min_size_test=800,
|
| 63 |
+
max_size_test=1333,
|
| 64 |
+
img_format="RGB",
|
| 65 |
+
metadata_dataset="coco_2017_val",
|
| 66 |
+
instance_mode=ColorMode.IMAGE,
|
| 67 |
+
parallel=False,
|
| 68 |
+
):
|
| 69 |
+
"""
|
| 70 |
+
Args:
|
| 71 |
+
cfg (CfgNode):
|
| 72 |
+
instance_mode (ColorMode):
|
| 73 |
+
parallel (bool): whether to run the model in different processes from visualization.
|
| 74 |
+
Useful since the visualization logic can be slow.
|
| 75 |
+
"""
|
| 76 |
+
self.metadata = MetadataCatalog.get(
|
| 77 |
+
metadata_dataset if metadata_dataset is not None else "__unused"
|
| 78 |
+
)
|
| 79 |
+
self.cpu_device = torch.device("cpu")
|
| 80 |
+
self.instance_mode = instance_mode
|
| 81 |
+
|
| 82 |
+
self.parallel = parallel
|
| 83 |
+
if parallel:
|
| 84 |
+
num_gpu = torch.cuda.device_count()
|
| 85 |
+
self.predictor = AsyncPredictor(
|
| 86 |
+
model=model,
|
| 87 |
+
min_size_test=min_size_test,
|
| 88 |
+
max_size_test=max_size_test,
|
| 89 |
+
img_format=img_format,
|
| 90 |
+
metadata_dataset=metadata_dataset,
|
| 91 |
+
num_gpus=num_gpu,
|
| 92 |
+
)
|
| 93 |
+
else:
|
| 94 |
+
self.predictor = DefaultPredictor(
|
| 95 |
+
model=model,
|
| 96 |
+
min_size_test=min_size_test,
|
| 97 |
+
max_size_test=max_size_test,
|
| 98 |
+
img_format=img_format,
|
| 99 |
+
metadata_dataset=metadata_dataset,
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
def run_on_image(self, image, threshold=0.5):
|
| 103 |
+
"""
|
| 104 |
+
Args:
|
| 105 |
+
image (np.ndarray): an image of shape (H, W, C) (in BGR order).
|
| 106 |
+
This is the format used by OpenCV.
|
| 107 |
+
|
| 108 |
+
Returns:
|
| 109 |
+
predictions (dict): the output of the model.
|
| 110 |
+
vis_output (VisImage): the visualized image output.
|
| 111 |
+
"""
|
| 112 |
+
vis_output = None
|
| 113 |
+
predictions = self.predictor(image)
|
| 114 |
+
predictions = filter_predictions_with_confidence(predictions, threshold)
|
| 115 |
+
# Convert image from OpenCV BGR format to Matplotlib RGB format.
|
| 116 |
+
image = image[:, :, ::-1]
|
| 117 |
+
visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode)
|
| 118 |
+
if "panoptic_seg" in predictions:
|
| 119 |
+
panoptic_seg, segments_info = predictions["panoptic_seg"]
|
| 120 |
+
vis_output = visualizer.draw_panoptic_seg_predictions(
|
| 121 |
+
panoptic_seg.to(self.cpu_device), segments_info
|
| 122 |
+
)
|
| 123 |
+
else:
|
| 124 |
+
if "sem_seg" in predictions:
|
| 125 |
+
vis_output = visualizer.draw_sem_seg(
|
| 126 |
+
predictions["sem_seg"].argmax(dim=0).to(self.cpu_device)
|
| 127 |
+
)
|
| 128 |
+
if "instances" in predictions:
|
| 129 |
+
instances = predictions["instances"].to(self.cpu_device)
|
| 130 |
+
vis_output = visualizer.draw_instance_predictions(predictions=instances)
|
| 131 |
+
|
| 132 |
+
return predictions, vis_output
|
| 133 |
+
|
| 134 |
+
def _frame_from_video(self, video):
|
| 135 |
+
while video.isOpened():
|
| 136 |
+
success, frame = video.read()
|
| 137 |
+
if success:
|
| 138 |
+
yield frame
|
| 139 |
+
else:
|
| 140 |
+
break
|
| 141 |
+
|
| 142 |
+
def run_on_video(self, video, threshold=0.5):
|
| 143 |
+
"""
|
| 144 |
+
Visualizes predictions on frames of the input video.
|
| 145 |
+
|
| 146 |
+
Args:
|
| 147 |
+
video (cv2.VideoCapture): a :class:`VideoCapture` object, whose source can be
|
| 148 |
+
either a webcam or a video file.
|
| 149 |
+
|
| 150 |
+
Yields:
|
| 151 |
+
ndarray: BGR visualizations of each video frame.
|
| 152 |
+
"""
|
| 153 |
+
video_visualizer = VideoVisualizer(self.metadata, self.instance_mode)
|
| 154 |
+
|
| 155 |
+
def process_predictions(frame, predictions, threshold):
|
| 156 |
+
predictions = filter_predictions_with_confidence(predictions, threshold)
|
| 157 |
+
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
| 158 |
+
if "panoptic_seg" in predictions:
|
| 159 |
+
panoptic_seg, segments_info = predictions["panoptic_seg"]
|
| 160 |
+
vis_frame = video_visualizer.draw_panoptic_seg_predictions(
|
| 161 |
+
frame, panoptic_seg.to(self.cpu_device), segments_info
|
| 162 |
+
)
|
| 163 |
+
elif "instances" in predictions:
|
| 164 |
+
predictions = predictions["instances"].to(self.cpu_device)
|
| 165 |
+
vis_frame = video_visualizer.draw_instance_predictions(frame, predictions)
|
| 166 |
+
elif "sem_seg" in predictions:
|
| 167 |
+
vis_frame = video_visualizer.draw_sem_seg(
|
| 168 |
+
frame, predictions["sem_seg"].argmax(dim=0).to(self.cpu_device)
|
| 169 |
+
)
|
| 170 |
+
|
| 171 |
+
# Converts Matplotlib RGB format to OpenCV BGR format
|
| 172 |
+
vis_frame = cv2.cvtColor(vis_frame.get_image(), cv2.COLOR_RGB2BGR)
|
| 173 |
+
return vis_frame
|
| 174 |
+
|
| 175 |
+
frame_gen = self._frame_from_video(video)
|
| 176 |
+
if self.parallel:
|
| 177 |
+
buffer_size = self.predictor.default_buffer_size
|
| 178 |
+
|
| 179 |
+
frame_data = deque()
|
| 180 |
+
|
| 181 |
+
for cnt, frame in enumerate(frame_gen):
|
| 182 |
+
frame_data.append(frame)
|
| 183 |
+
self.predictor.put(frame)
|
| 184 |
+
|
| 185 |
+
if cnt >= buffer_size:
|
| 186 |
+
frame = frame_data.popleft()
|
| 187 |
+
predictions = self.predictor.get()
|
| 188 |
+
yield process_predictions(frame, predictions, threshold)
|
| 189 |
+
|
| 190 |
+
while len(frame_data):
|
| 191 |
+
frame = frame_data.popleft()
|
| 192 |
+
predictions = self.predictor.get()
|
| 193 |
+
yield process_predictions(frame, predictions, threshold)
|
| 194 |
+
else:
|
| 195 |
+
for frame in frame_gen:
|
| 196 |
+
yield process_predictions(frame, self.predictor(frame), threshold)
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
class DefaultPredictor:
|
| 200 |
+
def __init__(
|
| 201 |
+
self,
|
| 202 |
+
model,
|
| 203 |
+
min_size_test=800,
|
| 204 |
+
max_size_test=1333,
|
| 205 |
+
img_format="RGB",
|
| 206 |
+
metadata_dataset="coco_2017_val",
|
| 207 |
+
):
|
| 208 |
+
self.model = model
|
| 209 |
+
# self.model.eval()
|
| 210 |
+
self.metadata = MetadataCatalog.get(metadata_dataset)
|
| 211 |
+
|
| 212 |
+
# checkpointer = DetectionCheckpointer(self.model)
|
| 213 |
+
# checkpointer.load(init_checkpoint)
|
| 214 |
+
|
| 215 |
+
self.aug = T.ResizeShortestEdge([min_size_test, min_size_test], max_size_test)
|
| 216 |
+
|
| 217 |
+
self.input_format = img_format
|
| 218 |
+
assert self.input_format in ["RGB", "BGR"], self.input_format
|
| 219 |
+
|
| 220 |
+
def __call__(self, original_image):
|
| 221 |
+
"""
|
| 222 |
+
Args:
|
| 223 |
+
original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).
|
| 224 |
+
|
| 225 |
+
Returns:
|
| 226 |
+
predictions (dict):
|
| 227 |
+
the output of the model for one image only.
|
| 228 |
+
See :doc:`/tutorials/models` for details about the format.
|
| 229 |
+
"""
|
| 230 |
+
with torch.no_grad(): # https://github.com/sphinx-doc/sphinx/issues/4258
|
| 231 |
+
# Apply pre-processing to image.
|
| 232 |
+
if self.input_format == "RGB":
|
| 233 |
+
# whether the model expects BGR inputs or RGB
|
| 234 |
+
original_image = original_image[:, :, ::-1]
|
| 235 |
+
height, width = original_image.shape[:2]
|
| 236 |
+
image = self.aug.get_transform(original_image).apply_image(original_image)
|
| 237 |
+
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
|
| 238 |
+
|
| 239 |
+
inputs = {"image": image, "height": height, "width": width}
|
| 240 |
+
predictions = self.model([inputs])[0]
|
| 241 |
+
return predictions
|
| 242 |
+
|
| 243 |
+
|
| 244 |
+
class AsyncPredictor:
|
| 245 |
+
"""
|
| 246 |
+
A predictor that runs the model asynchronously, possibly on >1 GPUs.
|
| 247 |
+
Because rendering the visualization takes considerably amount of time,
|
| 248 |
+
this helps improve throughput a little bit when rendering videos.
|
| 249 |
+
"""
|
| 250 |
+
|
| 251 |
+
class _StopToken:
|
| 252 |
+
pass
|
| 253 |
+
|
| 254 |
+
class _PredictWorker(mp.Process):
|
| 255 |
+
def __init__(
|
| 256 |
+
self,
|
| 257 |
+
model,
|
| 258 |
+
task_queue,
|
| 259 |
+
result_queue,
|
| 260 |
+
min_size_test=800,
|
| 261 |
+
max_size_test=1333,
|
| 262 |
+
img_format="RGB",
|
| 263 |
+
metadata_dataset="coco_2017_val",
|
| 264 |
+
):
|
| 265 |
+
self.model = model
|
| 266 |
+
self.min_size_test = min_size_test
|
| 267 |
+
self.max_size_test = max_size_test
|
| 268 |
+
self.img_format = img_format
|
| 269 |
+
self.metadata_dataset = metadata_dataset
|
| 270 |
+
self.task_queue = task_queue
|
| 271 |
+
self.result_queue = result_queue
|
| 272 |
+
super().__init__()
|
| 273 |
+
|
| 274 |
+
def run(self):
|
| 275 |
+
predictor = DefaultPredictor(
|
| 276 |
+
model=self.model,
|
| 277 |
+
min_size_test=self.min_size_test,
|
| 278 |
+
max_size_test=self.max_size_test,
|
| 279 |
+
img_format=self.img_format,
|
| 280 |
+
metadata_dataset=self.metadata_dataset,
|
| 281 |
+
)
|
| 282 |
+
|
| 283 |
+
while True:
|
| 284 |
+
task = self.task_queue.get()
|
| 285 |
+
if isinstance(task, AsyncPredictor._StopToken):
|
| 286 |
+
break
|
| 287 |
+
idx, data = task
|
| 288 |
+
result = predictor(data)
|
| 289 |
+
self.result_queue.put((idx, result))
|
| 290 |
+
|
| 291 |
+
def __init__(self, cfg, num_gpus: int = 1):
|
| 292 |
+
"""
|
| 293 |
+
Args:
|
| 294 |
+
cfg (CfgNode):
|
| 295 |
+
num_gpus (int): if 0, will run on CPU
|
| 296 |
+
"""
|
| 297 |
+
num_workers = max(num_gpus, 1)
|
| 298 |
+
self.task_queue = mp.Queue(maxsize=num_workers * 3)
|
| 299 |
+
self.result_queue = mp.Queue(maxsize=num_workers * 3)
|
| 300 |
+
self.procs = []
|
| 301 |
+
for gpuid in range(max(num_gpus, 1)):
|
| 302 |
+
cfg = cfg.clone()
|
| 303 |
+
cfg.defrost()
|
| 304 |
+
cfg.MODEL.DEVICE = "cuda:{}".format(gpuid) if num_gpus > 0 else "cpu"
|
| 305 |
+
self.procs.append(
|
| 306 |
+
AsyncPredictor._PredictWorker(cfg, self.task_queue, self.result_queue)
|
| 307 |
+
)
|
| 308 |
+
|
| 309 |
+
self.put_idx = 0
|
| 310 |
+
self.get_idx = 0
|
| 311 |
+
self.result_rank = []
|
| 312 |
+
self.result_data = []
|
| 313 |
+
|
| 314 |
+
for p in self.procs:
|
| 315 |
+
p.start()
|
| 316 |
+
atexit.register(self.shutdown)
|
| 317 |
+
|
| 318 |
+
def put(self, image):
|
| 319 |
+
self.put_idx += 1
|
| 320 |
+
self.task_queue.put((self.put_idx, image))
|
| 321 |
+
|
| 322 |
+
def get(self):
|
| 323 |
+
self.get_idx += 1 # the index needed for this request
|
| 324 |
+
if len(self.result_rank) and self.result_rank[0] == self.get_idx:
|
| 325 |
+
res = self.result_data[0]
|
| 326 |
+
del self.result_data[0], self.result_rank[0]
|
| 327 |
+
return res
|
| 328 |
+
|
| 329 |
+
while True:
|
| 330 |
+
# make sure the results are returned in the correct order
|
| 331 |
+
idx, res = self.result_queue.get()
|
| 332 |
+
if idx == self.get_idx:
|
| 333 |
+
return res
|
| 334 |
+
insert = bisect.bisect(self.result_rank, idx)
|
| 335 |
+
self.result_rank.insert(insert, idx)
|
| 336 |
+
self.result_data.insert(insert, res)
|
| 337 |
+
|
| 338 |
+
def __len__(self):
|
| 339 |
+
return self.put_idx - self.get_idx
|
| 340 |
+
|
| 341 |
+
def __call__(self, image):
|
| 342 |
+
self.put(image)
|
| 343 |
+
return self.get()
|
| 344 |
+
|
| 345 |
+
def shutdown(self):
|
| 346 |
+
for _ in self.procs:
|
| 347 |
+
self.task_queue.put(AsyncPredictor._StopToken())
|
| 348 |
+
|
| 349 |
+
@property
|
| 350 |
+
def default_buffer_size(self):
|
| 351 |
+
return len(self.procs) * 5
|
| 352 |
+
|
| 353 |
+
|
| 354 |
+
detrex_model_list = {
|
| 355 |
+
# DETR
|
| 356 |
+
"detr/detr_r50_300ep": {
|
| 357 |
+
"configs": "projects/detr/configs/detr_r50_300ep.py",
|
| 358 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/converted_detr_r50_500ep.pth"
|
| 359 |
+
},
|
| 360 |
+
"detr/detr_r50_dc5_300ep": {
|
| 361 |
+
"configs": "projects/detr/configs/detr_r50_dc5_300ep.py",
|
| 362 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_detr_r50_dc5.pth"
|
| 363 |
+
},
|
| 364 |
+
"detr/detr_r101_300ep.py": {
|
| 365 |
+
"configs": "projects/detr/configs/detr_r101_300ep.py",
|
| 366 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/converted_detr_r101_500ep.pth"
|
| 367 |
+
},
|
| 368 |
+
"detr/detr_r101_dc5_300ep.py": {
|
| 369 |
+
"configs": "projects/detr/configs/detr_r101_dc5_300ep.py",
|
| 370 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_detr_r101_dc5.pth"
|
| 371 |
+
},
|
| 372 |
+
# Deformable-DETR
|
| 373 |
+
"deformable_detr/deformable_detr_r50_with_box_refinement_50ep": {
|
| 374 |
+
"configs": "projects/deformable_detr/configs/deformable_detr_r50_with_box_refinement_50ep.py",
|
| 375 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/deformable_detr_with_box_refinement_50ep_new.pth"
|
| 376 |
+
},
|
| 377 |
+
"deformable_detr/deformable_detr_r50_two_stage_50ep": {
|
| 378 |
+
"configs": "projects/deformable_detr/configs/deformable_detr_r50_two_stage_50ep.py",
|
| 379 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/deformable_detr_r50_two_stage_50ep_new.pth"
|
| 380 |
+
},
|
| 381 |
+
# Anchor-DETR
|
| 382 |
+
"anchor_detr/anchor_detr_r50_50ep":{
|
| 383 |
+
"configs":"projects/anchor_detr/configs/anchor_detr_r50_50ep.py",
|
| 384 |
+
"ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/anchor_detr_r50_50ep.pth"
|
| 385 |
+
},
|
| 386 |
+
"anchor_detr/anchor_detr_r50_50ep_(converted)":{
|
| 387 |
+
"configs":"projects/anchor_detr/configs/anchor_detr_r50_50ep.py",
|
| 388 |
+
"ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_anchor_detr_r50_50ep.pth"
|
| 389 |
+
},
|
| 390 |
+
"anchor_detr/anchor_detr_r50_dc5_50ep":{
|
| 391 |
+
"configs":"projects/anchor_detr/configs/anchor_detr_r50_dc5_50ep.py",
|
| 392 |
+
"ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_anchor_detr_r50_dc5_50ep.pth"
|
| 393 |
+
},
|
| 394 |
+
"anchor_detr/anchor_detr_r101_50ep":{
|
| 395 |
+
"configs":"projects/anchor_detr/configs/anchor_detr_r101_dc5_50ep.py",
|
| 396 |
+
"ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_anchor_detr_r101_dc5_50ep.pth"
|
| 397 |
+
},
|
| 398 |
+
"anchor_detr/anchor_detr_r101_dc5_50ep":{
|
| 399 |
+
"configs":"projects/anchor_detr/configs/anchor_detr_r101_dc5_50ep.py",
|
| 400 |
+
"ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_anchor_detr_r101_50ep.pth"
|
| 401 |
+
},
|
| 402 |
+
# Conditional-DETR
|
| 403 |
+
"conditional_detr/conditional_detr_r50_50ep":{
|
| 404 |
+
"configs":"projects/conditional_detr/configs/conditional_detr_r50_50ep.py",
|
| 405 |
+
"ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/conditional_detr_r50_50ep.pth"
|
| 406 |
+
},
|
| 407 |
+
"conditional_detr/conditional_detr_r50_50ep_(converted)":{
|
| 408 |
+
"configs":"",
|
| 409 |
+
"ckpts":""
|
| 410 |
+
},
|
| 411 |
+
"conditional_detr/conditional_detr_r101_50ep":{
|
| 412 |
+
"configs":"projects/conditional_detr/configs/conditional_detr_r101_50ep.py",
|
| 413 |
+
"ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/converted_conditional_detr_r101_50ep.pth"
|
| 414 |
+
},
|
| 415 |
+
"conditional_detr/conditional_detr_r101_dc5_50ep": {
|
| 416 |
+
"configs": "projects/conditional_detr/configs/conditional_detr_r101_dc5_50ep.py",
|
| 417 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_conditional_detr_r101_dc5.pth"
|
| 418 |
+
},
|
| 419 |
+
# DAB-DETR
|
| 420 |
+
"dab_detr/dab_detr_r50_50ep": {
|
| 421 |
+
"configs": "projects/dab_detr/configs/dab_detr_r50_50ep.py",
|
| 422 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/dab_detr_r50_50ep.pth"
|
| 423 |
+
},
|
| 424 |
+
"dab_detr/dab_detr_r50_3patterns_50ep": {
|
| 425 |
+
"configs": "projects/dab_detr/configs/dab_detr_r50_3patterns_50ep.py",
|
| 426 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_dab_detr_r50_3patterns.pth"
|
| 427 |
+
},
|
| 428 |
+
"dab_detr/dab_detr_r50_dc5_50ep": {
|
| 429 |
+
"configs": "projects/dab_detr/configs/dab_detr_r50_dc5_50ep.py",
|
| 430 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_dab_detr_r50_dc5.pth"
|
| 431 |
+
},
|
| 432 |
+
"dab_detr/dab_detr_r50_dc5_3patterns_50ep": {
|
| 433 |
+
"configs": "projects/dab_detr/configs/dab_detr_r50_3patterns_50ep.py",
|
| 434 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_dab_detr_r50_dc5_3patterns.pth"
|
| 435 |
+
},
|
| 436 |
+
"dab_detr/dab_detr_r101_50ep": {
|
| 437 |
+
"configs": "projects/dab_detr/configs/dab_detr_r101_50ep.py",
|
| 438 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/dab_detr_r101_50ep.pth"
|
| 439 |
+
},
|
| 440 |
+
"dab_detr/dab_detr_r50_dc5_3patterns_50ep_(converted)": {
|
| 441 |
+
"configs": "projects/dab_detr/configs/dab_detr_r50_dc5_3patterns_50ep.py",
|
| 442 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_detr_r101_dc5.pth"
|
| 443 |
+
},
|
| 444 |
+
"dab_detr/dab_detr_swin_t_in1k_50ep": {
|
| 445 |
+
"configs": "projects/dab_detr/configs/dab_detr_swin_t_in1k_50ep.py",
|
| 446 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/dab_detr_swin_t_in1k_50ep.pth"
|
| 447 |
+
},
|
| 448 |
+
"dab_detr/dab_deformable_detr_r50_50ep": {
|
| 449 |
+
"configs": "projects/dab_detr/configs/dab_deformable_detr_r50_50ep.py",
|
| 450 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dab_deformable_detr_r50_50ep_49AP.pth"
|
| 451 |
+
},
|
| 452 |
+
"dab_detr/dab_deformable_detr_r50_two_stage_50ep": {
|
| 453 |
+
"configs": "projects/dab_detr/configs/dab_deformable_detr_r50_two_stage_50ep.py",
|
| 454 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dab_deformable_detr_r50_two_stage_49_7AP.pth"
|
| 455 |
+
},
|
| 456 |
+
# DN-DETR
|
| 457 |
+
"dn_detr/dn_detr_r50_50ep": {
|
| 458 |
+
"configs": "projects/dn_detr/configs/dn_detr_r50_50ep.py",
|
| 459 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/dn_detr_r50_50ep.pth"
|
| 460 |
+
},
|
| 461 |
+
"dn_detr/dn_detr_r50_dc5_50ep": {
|
| 462 |
+
"configs": "projects/dn_detr/configs/dn_detr_r50_dc5_50ep.py",
|
| 463 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_dn_detr_r50_dc5.pth"
|
| 464 |
+
},
|
| 465 |
+
# DINO
|
| 466 |
+
"dino/dino-resnet/dino_r50_5scale_12ep": {
|
| 467 |
+
"configs": "projects/dino/configs/dino-resnet/dino_r50_5scale_12ep.py",
|
| 468 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_r50_5scale_12ep.pth"
|
| 469 |
+
},
|
| 470 |
+
"dino/dino-resnet/dino_r50_4scale_12ep_300dn": {
|
| 471 |
+
"configs": "projects/dino/configs/dino-resnet/dino_r50_4scale_12ep_300dn.py",
|
| 472 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dino_r50_4scale_12ep_300dn.pth"
|
| 473 |
+
},
|
| 474 |
+
"dino/dino-resnet/dino_r50_4scale_24ep": {
|
| 475 |
+
"configs": "projects/dino/configs/dino-resnet/dino_r50_4scale_24ep.py",
|
| 476 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_r50_4scale_24ep.pth"
|
| 477 |
+
},
|
| 478 |
+
"dino/dino-resnet/dino_r101_4scale_12ep_": {
|
| 479 |
+
"configs": "projects/dino/configs/dino-resnet/dino_r101_4scale_12ep.py",
|
| 480 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_r101_4scale_12ep.pth"
|
| 481 |
+
},
|
| 482 |
+
# Pretrained DINO with Swin-Transformer Backbone
|
| 483 |
+
"dino/dino-swin/dino_swin_tiny_224_4scale_12ep": {
|
| 484 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_tiny_224_4scale_12ep.py",
|
| 485 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_tiny_224_4scale_12ep.pth"
|
| 486 |
+
},
|
| 487 |
+
"dino/dino-swin/dino_swin_tiny_224_22kto1k_finetune_4scale_12ep": {
|
| 488 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_tiny_224_4scale_12ep.py",
|
| 489 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_tiny_224_22kto1k_finetune_4scale_12ep.pth"
|
| 490 |
+
},
|
| 491 |
+
"dino/dino-swin/dino_swin_small_224_4scale_12ep": {
|
| 492 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_tiny_224_4scale_12ep.py",
|
| 493 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_small_224_4scale_12ep.pth"
|
| 494 |
+
},
|
| 495 |
+
"dino/dino-swin/dino_swin_base_384_4scale_12ep": {
|
| 496 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_base_384_4scale_12ep.py",
|
| 497 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_base_384_4scale_12ep.pth"
|
| 498 |
+
},
|
| 499 |
+
"dino/dino-swin/dino_swin_large_224_4scale_12ep": {
|
| 500 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_large_224_4scale_12ep.py",
|
| 501 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_large_224_4scale_12ep.pth"
|
| 502 |
+
},
|
| 503 |
+
"dino/dino-swin/dino_swin_large_384_4scale_12ep": {
|
| 504 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_large_384_4scale_12ep.py",
|
| 505 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_large_4scale_12ep.pth"
|
| 506 |
+
},
|
| 507 |
+
"dino/dino-swin/dino_swin_large_384_5scale_12ep": {
|
| 508 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_large_384_5scale_12ep.py",
|
| 509 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_swin_large_384_5scale_12ep.pth"
|
| 510 |
+
},
|
| 511 |
+
"dino/dino-swin/dino_swin_large_384_4scale_36ep": {
|
| 512 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_large_384_4scale_36ep.py",
|
| 513 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dino_swin_large_384_4scale_36ep.pth"
|
| 514 |
+
},
|
| 515 |
+
"dino/dino-swin/dino_swin_large_384_5scale_36ep": {
|
| 516 |
+
"configs": "projects/dino/configs/dino-swin/dino_swin_large_384_5scale_36ep.py",
|
| 517 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_swin_large_384_5scale_36ep.pth"
|
| 518 |
+
},
|
| 519 |
+
# Pretrained DINO with FocalNet Backbone
|
| 520 |
+
"dino/dino-swin/dino_focalnet_large_lrf_384_4scale_12ep": {
|
| 521 |
+
"configs": "projects/dino/configs/dino-focal/dino_focalnet_large_lrf_384_4scale_12ep.py",
|
| 522 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_focal_large_lrf_384_4scale_12ep.pth"
|
| 523 |
+
},
|
| 524 |
+
"dino/dino-swin/dino_focalnet_large_lrf_384_fl4_4scale_12ep": {
|
| 525 |
+
"configs": "projects/dino/configs/dino-focal/dino_focalnet_large_lrf_384_fl4_4scale_12ep.py",
|
| 526 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_focal_large_lrf_384_4scale_12ep.pth"
|
| 527 |
+
},
|
| 528 |
+
"dino/dino-swin/dino_focalnet_large_lrf_384_fl4_5scale_12ep": {
|
| 529 |
+
"configs": "projects/dino/configs/dino-focal/dino_focalnet_large_lrf_384_fl4_5scale_12ep.py",
|
| 530 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_focalnet_large_lrf_384_fl4_5scale_12ep.pth"
|
| 531 |
+
},
|
| 532 |
+
# Pretrained DINO with ViTDet Backbone
|
| 533 |
+
"dino/dino-vitdet/dino_vitdet_base_4scale_12ep": {
|
| 534 |
+
"configs": "projects/dino/configs/dino-vitdet/dino_vitdet_base_4scale_12ep.py",
|
| 535 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_vitdet_4scale_12ep.pth"
|
| 536 |
+
},
|
| 537 |
+
"dino/dino-vitdet/dino_vitdet_base_4scale_50ep": {
|
| 538 |
+
"configs": "projects/dino/configs/dino-vitdet/dino_vitdet_base_4scale_50ep.py",
|
| 539 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_vitdet_base_4scale_50ep.pth"
|
| 540 |
+
},
|
| 541 |
+
"dino/dino-vitdet/dino_vitdet_large_4scale_12ep": {
|
| 542 |
+
"configs": "projects/dino/configs/dino-vitdet/dino_vitdet_large_4scale_12ep.py",
|
| 543 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_vitdet_large_4scale_12ep.pth"
|
| 544 |
+
},
|
| 545 |
+
"dino/dino-vitdet/dino_vitdet_large_4scale_50ep": {
|
| 546 |
+
"configs": "projects/dino/configs/dino-vitdet/dino_vitdet_large_4scale_50ep.py",
|
| 547 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_vitdet_large_4scale_50ep.pth"
|
| 548 |
+
},
|
| 549 |
+
# H-Deformable-DETR
|
| 550 |
+
"h_deformable_detr/h_deformable_detr_r50_two_stage_12ep": {
|
| 551 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_r50_two_stage_12ep.py",
|
| 552 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.4.0/h_deformable_detr_r50_two_stage_12ep_modified_train_net.pth"
|
| 553 |
+
},
|
| 554 |
+
"h_deformable_detr/h_deformable_detr_r50_two_stage_12ep(converted)": {
|
| 555 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_r50_two_stage_12ep.py",
|
| 556 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/r50_hybrid_branch_lambda1_group6_t1500_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_12eps.pth"
|
| 557 |
+
},
|
| 558 |
+
"h_deformable_detr/h_deformable_detr_r50_two_stage_36ep": {
|
| 559 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_r50_two_stage_36ep.py",
|
| 560 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/r50_hybrid_branch_lambda1_group6_t1500_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_36eps.pth"
|
| 561 |
+
},
|
| 562 |
+
"h_deformable_detr/h_deformable_detr_swin_tiny_two_stage_12ep": {
|
| 563 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_tiny_two_stage_12ep.py",
|
| 564 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/swin_tiny_hybrid_branch_lambda1_group6_t1500_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_12eps.pth"
|
| 565 |
+
},
|
| 566 |
+
"h_deformable_detr/h_deformable_detr_swin_tiny_two_stage_36ep": {
|
| 567 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_tiny_two_stage_36ep.py",
|
| 568 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/swin_tiny_hybrid_branch_lambda1_group6_t1500_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_36eps.pth"
|
| 569 |
+
},
|
| 570 |
+
"h_deformable_detr/h_deformable_detr_swin_large_two_stage_12ep": {
|
| 571 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_large_two_stage_12ep.py",
|
| 572 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/swin_large_hybrid_branch_lambda1_group6_t1500_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_12eps.pth"
|
| 573 |
+
},
|
| 574 |
+
"h_deformable_detr/h_deformable_detr_swin_large_two_stage_36ep": {
|
| 575 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_large_two_stage_36ep.py",
|
| 576 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/drop_path0.5_swin_large_hybrid_branch_lambda1_group6_t1500_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_36eps.pth"
|
| 577 |
+
},
|
| 578 |
+
"h_deformable_detr/h_deformable_detr_swin_large_two_stage_12ep_900queries": {
|
| 579 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_large_two_stage_12ep_900queries.py",
|
| 580 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/swin_large_hybrid_branch_lambda1_group6_t1500_n900_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_12eps.pth"
|
| 581 |
+
},
|
| 582 |
+
"h_deformable_detr/h_deformable_detr_swin_large_two_stage_36ep_900queries": {
|
| 583 |
+
"configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_large_two_stage_36ep_900queries.py",
|
| 584 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/drop_path0.5_swin_large_hybrid_branch_lambda1_group6_t1500_n900_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_36eps.pth"
|
| 585 |
+
},
|
| 586 |
+
# DETA
|
| 587 |
+
"deta/improved_deformable_detr_baseline_50ep": {
|
| 588 |
+
"configs": "projects/deta/configs/improved_deformable_detr_baseline_50ep.py",
|
| 589 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_deta_improved_deformable_baseline.pth"
|
| 590 |
+
},
|
| 591 |
+
"deta/deta_r50_5scale_12ep_bs8": {
|
| 592 |
+
"configs": "projects/deta/configs/deta_r50_5scale_12ep_bs8.py",
|
| 593 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.4.0/deta_r50_5scale_12ep_bs8.pth"
|
| 594 |
+
},
|
| 595 |
+
"deta/deta_r50_5scale_12ep": {
|
| 596 |
+
"configs": "projects/deta/configs/deta_r50_5scale_12ep.py",
|
| 597 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.4.0/deta_r50_5scale_12ep_hacked_trainer.pth"
|
| 598 |
+
},
|
| 599 |
+
"deta/deta_r50_5scale_no_frozen_backbone": {
|
| 600 |
+
"configs": "projects/deta/configs/deta_r50_5scale_no_frozen_backbone.py",
|
| 601 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.4.0/deta_r50_5scale_12ep_no_freeze_backbone.pth"
|
| 602 |
+
},
|
| 603 |
+
"deta/deta_r50_5scale_12ep(converted)": {
|
| 604 |
+
"configs": "projects/deta/configs/deta_r50_5scale_12ep.py",
|
| 605 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_deta_r50_5scale_12ep.pth"
|
| 606 |
+
},
|
| 607 |
+
"deta/DETA-Swin-Large-finetune (converted)": {
|
| 608 |
+
"configs": "projects/deta/configs/deta_r50_5scale_12ep.py",
|
| 609 |
+
"ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_deta_swin_o365_finetune.pth"
|
| 610 |
+
},
|
| 611 |
+
}
|
| 612 |
+
|
| 613 |
+
|
| 614 |
+
def setup(args):
|
| 615 |
+
cfg = LazyConfig.load(args.config_file)
|
| 616 |
+
cfg = LazyConfig.apply_overrides(cfg, args.opts)
|
| 617 |
+
return cfg
|
| 618 |
+
|
| 619 |
+
|
| 620 |
+
def get_parser():
|
| 621 |
+
parser = argparse.ArgumentParser(description="detrex demo for visualizing customized inputs")
|
| 622 |
+
parser.add_argument(
|
| 623 |
+
"--config-file",
|
| 624 |
+
default="projects/dino/configs/dino_r50_4scale_12ep.py",
|
| 625 |
+
metavar="FILE",
|
| 626 |
+
help="path to config file",
|
| 627 |
+
)
|
| 628 |
+
parser.add_argument("--webcam", action="store_true", help="Take inputs from webcam.")
|
| 629 |
+
parser.add_argument("--video-input", help="Path to video file.")
|
| 630 |
+
parser.add_argument(
|
| 631 |
+
"--input",
|
| 632 |
+
nargs="+",
|
| 633 |
+
help="A list of space separated input images; "
|
| 634 |
+
"or a single glob pattern such as 'directory/*.jpg'",
|
| 635 |
+
)
|
| 636 |
+
parser.add_argument(
|
| 637 |
+
"--output",
|
| 638 |
+
help="A file or directory to save output visualizations. "
|
| 639 |
+
"If not given, will show output in an OpenCV window.",
|
| 640 |
+
)
|
| 641 |
+
parser.add_argument(
|
| 642 |
+
"--min_size_test",
|
| 643 |
+
type=int,
|
| 644 |
+
default=800,
|
| 645 |
+
help="Size of the smallest side of the image during testing. Set to zero to disable resize in testing.",
|
| 646 |
+
)
|
| 647 |
+
parser.add_argument(
|
| 648 |
+
"--max_size_test",
|
| 649 |
+
type=float,
|
| 650 |
+
default=1333,
|
| 651 |
+
help="Maximum size of the side of the image during testing.",
|
| 652 |
+
)
|
| 653 |
+
parser.add_argument(
|
| 654 |
+
"--img_format",
|
| 655 |
+
type=str,
|
| 656 |
+
default="RGB",
|
| 657 |
+
help="The format of the loading images.",
|
| 658 |
+
)
|
| 659 |
+
parser.add_argument(
|
| 660 |
+
"--metadata_dataset",
|
| 661 |
+
type=str,
|
| 662 |
+
default="coco_2017_val",
|
| 663 |
+
help="The metadata infomation to be used. Default to COCO val metadata.",
|
| 664 |
+
)
|
| 665 |
+
parser.add_argument(
|
| 666 |
+
"--confidence-threshold",
|
| 667 |
+
type=float,
|
| 668 |
+
default=0.5,
|
| 669 |
+
help="Minimum score for instance predictions to be shown",
|
| 670 |
+
)
|
| 671 |
+
parser.add_argument(
|
| 672 |
+
"--opts",
|
| 673 |
+
help="Modify config options using the command-line",
|
| 674 |
+
default=None,
|
| 675 |
+
nargs=argparse.REMAINDER,
|
| 676 |
+
)
|
| 677 |
+
return parser
|
| 678 |
+
|
| 679 |
+
|
| 680 |
+
def test_opencv_video_format(codec, file_ext):
|
| 681 |
+
with tempfile.TemporaryDirectory(prefix="video_format_test") as dir:
|
| 682 |
+
filename = os.path.join(dir, "test_file" + file_ext)
|
| 683 |
+
writer = cv2.VideoWriter(
|
| 684 |
+
filename=filename,
|
| 685 |
+
fourcc=cv2.VideoWriter_fourcc(*codec),
|
| 686 |
+
fps=float(30),
|
| 687 |
+
frameSize=(10, 10),
|
| 688 |
+
isColor=True,
|
| 689 |
+
)
|
| 690 |
+
[writer.write(np.zeros((10, 10, 3), np.uint8)) for _ in range(30)]
|
| 691 |
+
writer.release()
|
| 692 |
+
if os.path.isfile(filename):
|
| 693 |
+
return True
|
| 694 |
+
return False
|
| 695 |
+
|
| 696 |
+
|
| 697 |
+
def download_ckpts_and_image(ckpts):
|
| 698 |
+
print("ckpts:", ckpts)
|
| 699 |
+
torch.hub.download_url_to_file(ckpts, "dino_deitsmall16_pretrain.pth")
|
| 700 |
+
|
| 701 |
+
def run_detection(input_file, output_file, model_name, input_confidence, device):
|
| 702 |
+
|
| 703 |
+
configs = detrex_model_list[model_name]["configs"]
|
| 704 |
+
ckpts = detrex_model_list[model_name]["ckpts"]
|
| 705 |
+
|
| 706 |
+
mp.set_start_method("spawn", force=True)
|
| 707 |
+
args = get_parser().parse_args([
|
| 708 |
+
"--config-file", configs,
|
| 709 |
+
"--input", input_file,
|
| 710 |
+
"--output", output_file,
|
| 711 |
+
"--confidence-threshold", str(input_confidence),
|
| 712 |
+
"--opts", "train.init_checkpoint=" + ckpts
|
| 713 |
+
])
|
| 714 |
+
setup_logger(name="fvcore")
|
| 715 |
+
logger = setup_logger()
|
| 716 |
+
logger.info("Arguments: " + str(args))
|
| 717 |
+
|
| 718 |
+
cfg = setup(args)
|
| 719 |
+
cfg.model.device = device
|
| 720 |
+
cfg.train.device = device
|
| 721 |
+
model = instantiate(cfg.model)
|
| 722 |
+
model.to(cfg.train.device)
|
| 723 |
+
checkpointer = DetectionCheckpointer(model)
|
| 724 |
+
checkpointer.load(cfg.train.init_checkpoint)
|
| 725 |
+
|
| 726 |
+
model.eval()
|
| 727 |
+
|
| 728 |
+
demo = VisualizationDemo(
|
| 729 |
+
model=model,
|
| 730 |
+
min_size_test=args.min_size_test,
|
| 731 |
+
max_size_test=args.max_size_test,
|
| 732 |
+
img_format=args.img_format,
|
| 733 |
+
metadata_dataset=args.metadata_dataset,
|
| 734 |
+
)
|
| 735 |
+
|
| 736 |
+
if args.input:
|
| 737 |
+
if len(args.input) == 1:
|
| 738 |
+
args.input = glob.glob(os.path.expanduser(args.input[0]))
|
| 739 |
+
assert args.input, "The input path(s) was not found"
|
| 740 |
+
for path in tqdm.tqdm(args.input, disable=not args.output):
|
| 741 |
+
# use PIL, to be consistent with evaluation
|
| 742 |
+
img = read_image(path, format="BGR")
|
| 743 |
+
start_time = time.time()
|
| 744 |
+
predictions, visualized_output = demo.run_on_image(img, args.confidence_threshold)
|
| 745 |
+
logger.info(
|
| 746 |
+
"{}: {} in {:.2f}s".format(
|
| 747 |
+
path,
|
| 748 |
+
"detected {} instances".format(len(predictions["instances"]))
|
| 749 |
+
if "instances" in predictions
|
| 750 |
+
else "finished",
|
| 751 |
+
time.time() - start_time,
|
| 752 |
+
)
|
| 753 |
+
)
|
| 754 |
+
|
| 755 |
+
if args.output:
|
| 756 |
+
if os.path.isdir(args.output):
|
| 757 |
+
assert os.path.isdir(args.output), args.output
|
| 758 |
+
out_filename = os.path.join(args.output, os.path.basename(path))
|
| 759 |
+
else:
|
| 760 |
+
assert len(args.input) == 1, "Please specify a directory with args.output"
|
| 761 |
+
out_filename = args.output
|
| 762 |
+
visualized_output.save(out_filename)
|
| 763 |
+
|
| 764 |
+
def download_test_img():
|
| 765 |
+
import shutil
|
| 766 |
+
torch.hub.download_url_to_file(
|
| 767 |
+
'https://github.com/isLinXu/issues/files/12658779/projects.zip',
|
| 768 |
+
'projects.zip')
|
| 769 |
+
# Images
|
| 770 |
+
torch.hub.download_url_to_file(
|
| 771 |
+
'https://user-images.githubusercontent.com/59380685/268517006-d8d4d3b3-964a-4f4d-8458-18c7eb75a4f2.jpg',
|
| 772 |
+
'000000502136.jpg')
|
| 773 |
+
shutil.unpack_archive('projects.zip', './', 'zip')
|
| 774 |
+
|
| 775 |
+
def detect_image(input_image, model_name, input_confidence, device):
|
| 776 |
+
input_dir = "input.jpg"
|
| 777 |
+
input_image.save(input_dir)
|
| 778 |
+
output_image = "output.jpg"
|
| 779 |
+
run_detection(input_dir, output_image, model_name, input_confidence, device)
|
| 780 |
+
return output_image
|
| 781 |
+
|
| 782 |
+
|
| 783 |
+
if __name__ == '__main__':
|
| 784 |
+
input_image = gr.inputs.Image(type='pil', label="Input Image")
|
| 785 |
+
input_model_name = gr.inputs.Dropdown(list(detrex_model_list.keys()), label="Model Name", default="dab_detr/dab_detr_r50_50ep")
|
| 786 |
+
input_confidence = gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.25, label="Confidence Threshold")
|
| 787 |
+
input_device = gr.inputs.Radio(["cpu", "cuda"], label="Device", default="cpu")
|
| 788 |
+
output_image = gr.outputs.Image(type='pil', label="Output Image")
|
| 789 |
+
download_test_img()
|
| 790 |
+
examples = [["000000502136.jpg", "dab_detr/dab_detr_r50_50ep", 0.25, "cpu"]]
|
| 791 |
+
title = "🦖detrex: Benchmarking Detection Transformers web demo"
|
| 792 |
+
description = "<div align='center'><img src='https://raw.githubusercontent.com/IDEA-Research/detrex/main/assets/logo_2.png' width='250''/><div>" \
|
| 793 |
+
"<p style='text-align: center'><a href='https://github.com/IDEA-Research/detrex'>detrex</a> detrex detrex 是一个开源工具箱,提供最先进的基于 Transformer 的检测算法。它建立在Detectron2之上,其模块设计部分借鉴了MMDetection和DETR。非常感谢他们组织良好的代码。主分支适用于Pytorch 1.10+或更高版本(我们推荐Pytorch 1.12)。" \
|
| 794 |
+
"detrex is a research platform for DETR-based object detection, segmentation, pose estimation and other visual recognition tasks.</p>"
|
| 795 |
+
article = "<p style='text-align: center'><a href='https://github.com/IDEA-Research/detrex'>detrex</a></p>" \
|
| 796 |
+
"<p style='text-align: center'><a href='https://github.com/isLinXu'>gradio build by gatilin</a></a></p>"
|
| 797 |
+
|
| 798 |
+
image_interface = gr.Interface(detect_image,
|
| 799 |
+
inputs=[input_image, input_model_name, input_confidence, input_device],
|
| 800 |
+
outputs=output_image,examples=examples,
|
| 801 |
+
title=title, article=article, description=description)
|
| 802 |
+
image_interface.launch()
|
requirements.txt
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
git+https://github.com/facebookresearch/detectron2.git@v0.6#egg=detectron2
|
| 2 |
+
ultralytics~=8.0.169
|
| 3 |
+
wget~=3.2
|
| 4 |
+
opencv-python~=4.6.0.66
|
| 5 |
+
numpy~=1.23.0
|
| 6 |
+
pillow~=9.4.0
|
| 7 |
+
gradio~=3.42.0
|
| 8 |
+
pyyaml~=6.0
|
| 9 |
+
wandb~=0.13.11
|
| 10 |
+
tqdm~=4.65.0
|
| 11 |
+
matplotlib~=3.7.1
|
| 12 |
+
pandas~=2.0.0
|
| 13 |
+
seaborn~=0.12.2
|
| 14 |
+
requests~=2.31.0
|
| 15 |
+
psutil~=5.9.4
|
| 16 |
+
thop~=0.1.1-2209072238
|
| 17 |
+
timm~=0.9.2
|
| 18 |
+
super-gradients~=3.2.0
|
| 19 |
+
openmim
|