Spaces:
Sleeping
Sleeping
| import os | |
| os.system("pip install torch torchvision") | |
| os.system("git clone https://github.com/IDEA-Research/detrex.git") | |
| os.system("python -m pip install git+https://github.com/facebookresearch/detectron2.git@v0.6#egg=detectron2") | |
| os.system("python -m pip install git+https://github.com/IDEA-Research/detrex.git@v0.5.0#egg=detrex") | |
| os.system("git submodule sync") | |
| os.system("git submodule update --init") | |
| os.system("pip install fairscale") | |
| os.system("pip install Pillow==9.5.0") | |
| # os.system("cd detrex && pip install -e .") | |
| import argparse | |
| import glob | |
| import multiprocessing as mp | |
| import numpy as np | |
| import os | |
| import sys | |
| import tempfile | |
| import time | |
| import warnings | |
| import cv2 | |
| import torch | |
| import tqdm | |
| import gradio as gr | |
| # from demo.predictors import VisualizationDemo | |
| from detectron2.checkpoint import DetectionCheckpointer | |
| from detectron2.config import LazyConfig, instantiate | |
| from detectron2.data.detection_utils import read_image | |
| from detectron2.utils.logger import setup_logger | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| import atexit | |
| import bisect | |
| from copy import copy | |
| import multiprocessing as mp | |
| from collections import deque | |
| import cv2 | |
| import torch | |
| import detectron2.data.transforms as T | |
| from detectron2.data import MetadataCatalog | |
| from detectron2.structures import Instances | |
| from detectron2.utils.video_visualizer import VideoVisualizer | |
| from detectron2.utils.visualizer import ColorMode, Visualizer | |
| def filter_predictions_with_confidence(predictions, confidence_threshold=0.5): | |
| if "instances" in predictions: | |
| preds = predictions["instances"] | |
| keep_idxs = preds.scores > confidence_threshold | |
| predictions = copy(predictions) # don't modify the original | |
| predictions["instances"] = preds[keep_idxs] | |
| return predictions | |
| class VisualizationDemo(object): | |
| def __init__( | |
| self, | |
| model, | |
| min_size_test=800, | |
| max_size_test=1333, | |
| img_format="RGB", | |
| metadata_dataset="coco_2017_val", | |
| instance_mode=ColorMode.IMAGE, | |
| parallel=False, | |
| ): | |
| """ | |
| Args: | |
| cfg (CfgNode): | |
| instance_mode (ColorMode): | |
| parallel (bool): whether to run the model in different processes from visualization. | |
| Useful since the visualization logic can be slow. | |
| """ | |
| self.metadata = MetadataCatalog.get( | |
| metadata_dataset if metadata_dataset is not None else "__unused" | |
| ) | |
| self.cpu_device = torch.device("cpu") | |
| self.instance_mode = instance_mode | |
| self.parallel = parallel | |
| if parallel: | |
| num_gpu = torch.cuda.device_count() | |
| self.predictor = AsyncPredictor( | |
| model=model, | |
| min_size_test=min_size_test, | |
| max_size_test=max_size_test, | |
| img_format=img_format, | |
| metadata_dataset=metadata_dataset, | |
| num_gpus=num_gpu, | |
| ) | |
| else: | |
| self.predictor = DefaultPredictor( | |
| model=model, | |
| min_size_test=min_size_test, | |
| max_size_test=max_size_test, | |
| img_format=img_format, | |
| metadata_dataset=metadata_dataset, | |
| ) | |
| def run_on_image(self, image, threshold=0.5): | |
| """ | |
| Args: | |
| image (np.ndarray): an image of shape (H, W, C) (in BGR order). | |
| This is the format used by OpenCV. | |
| Returns: | |
| predictions (dict): the output of the model. | |
| vis_output (VisImage): the visualized image output. | |
| """ | |
| vis_output = None | |
| predictions = self.predictor(image) | |
| predictions = filter_predictions_with_confidence(predictions, threshold) | |
| # Convert image from OpenCV BGR format to Matplotlib RGB format. | |
| image = image[:, :, ::-1] | |
| visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode) | |
| if "panoptic_seg" in predictions: | |
| panoptic_seg, segments_info = predictions["panoptic_seg"] | |
| vis_output = visualizer.draw_panoptic_seg_predictions( | |
| panoptic_seg.to(self.cpu_device), segments_info | |
| ) | |
| else: | |
| if "sem_seg" in predictions: | |
| vis_output = visualizer.draw_sem_seg( | |
| predictions["sem_seg"].argmax(dim=0).to(self.cpu_device) | |
| ) | |
| if "instances" in predictions: | |
| instances = predictions["instances"].to(self.cpu_device) | |
| vis_output = visualizer.draw_instance_predictions(predictions=instances) | |
| return predictions, vis_output | |
| def _frame_from_video(self, video): | |
| while video.isOpened(): | |
| success, frame = video.read() | |
| if success: | |
| yield frame | |
| else: | |
| break | |
| def run_on_video(self, video, threshold=0.5): | |
| """ | |
| Visualizes predictions on frames of the input video. | |
| Args: | |
| video (cv2.VideoCapture): a :class:`VideoCapture` object, whose source can be | |
| either a webcam or a video file. | |
| Yields: | |
| ndarray: BGR visualizations of each video frame. | |
| """ | |
| video_visualizer = VideoVisualizer(self.metadata, self.instance_mode) | |
| def process_predictions(frame, predictions, threshold): | |
| predictions = filter_predictions_with_confidence(predictions, threshold) | |
| 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(self.cpu_device), segments_info | |
| ) | |
| elif "instances" in predictions: | |
| predictions = predictions["instances"].to(self.cpu_device) | |
| 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(self.cpu_device) | |
| ) | |
| # Converts Matplotlib RGB format to OpenCV BGR format | |
| vis_frame = cv2.cvtColor(vis_frame.get_image(), cv2.COLOR_RGB2BGR) | |
| return vis_frame | |
| frame_gen = self._frame_from_video(video) | |
| if self.parallel: | |
| buffer_size = self.predictor.default_buffer_size | |
| frame_data = deque() | |
| for cnt, frame in enumerate(frame_gen): | |
| frame_data.append(frame) | |
| self.predictor.put(frame) | |
| if cnt >= buffer_size: | |
| frame = frame_data.popleft() | |
| predictions = self.predictor.get() | |
| yield process_predictions(frame, predictions, threshold) | |
| while len(frame_data): | |
| frame = frame_data.popleft() | |
| predictions = self.predictor.get() | |
| yield process_predictions(frame, predictions, threshold) | |
| else: | |
| for frame in frame_gen: | |
| yield process_predictions(frame, self.predictor(frame), threshold) | |
| class DefaultPredictor: | |
| def __init__( | |
| self, | |
| model, | |
| min_size_test=800, | |
| max_size_test=1333, | |
| img_format="RGB", | |
| metadata_dataset="coco_2017_val", | |
| ): | |
| self.model = model | |
| # self.model.eval() | |
| self.metadata = MetadataCatalog.get(metadata_dataset) | |
| # checkpointer = DetectionCheckpointer(self.model) | |
| # checkpointer.load(init_checkpoint) | |
| self.aug = T.ResizeShortestEdge([min_size_test, min_size_test], max_size_test) | |
| self.input_format = img_format | |
| assert self.input_format in ["RGB", "BGR"], self.input_format | |
| def __call__(self, original_image): | |
| """ | |
| Args: | |
| original_image (np.ndarray): an image of shape (H, W, C) (in BGR order). | |
| Returns: | |
| predictions (dict): | |
| the output of the model for one image only. | |
| See :doc:`/tutorials/models` for details about the format. | |
| """ | |
| with torch.no_grad(): # https://github.com/sphinx-doc/sphinx/issues/4258 | |
| # Apply pre-processing to image. | |
| if self.input_format == "RGB": | |
| # whether the model expects BGR inputs or RGB | |
| original_image = original_image[:, :, ::-1] | |
| height, width = original_image.shape[:2] | |
| image = self.aug.get_transform(original_image).apply_image(original_image) | |
| image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1)) | |
| inputs = {"image": image, "height": height, "width": width} | |
| predictions = self.model([inputs])[0] | |
| return predictions | |
| class AsyncPredictor: | |
| """ | |
| A predictor that runs the model asynchronously, possibly on >1 GPUs. | |
| Because rendering the visualization takes considerably amount of time, | |
| this helps improve throughput a little bit when rendering videos. | |
| """ | |
| class _StopToken: | |
| pass | |
| class _PredictWorker(mp.Process): | |
| def __init__( | |
| self, | |
| model, | |
| task_queue, | |
| result_queue, | |
| min_size_test=800, | |
| max_size_test=1333, | |
| img_format="RGB", | |
| metadata_dataset="coco_2017_val", | |
| ): | |
| self.model = model | |
| self.min_size_test = min_size_test | |
| self.max_size_test = max_size_test | |
| self.img_format = img_format | |
| self.metadata_dataset = metadata_dataset | |
| self.task_queue = task_queue | |
| self.result_queue = result_queue | |
| super().__init__() | |
| def run(self): | |
| predictor = DefaultPredictor( | |
| model=self.model, | |
| min_size_test=self.min_size_test, | |
| max_size_test=self.max_size_test, | |
| img_format=self.img_format, | |
| metadata_dataset=self.metadata_dataset, | |
| ) | |
| while True: | |
| task = self.task_queue.get() | |
| if isinstance(task, AsyncPredictor._StopToken): | |
| break | |
| idx, data = task | |
| result = predictor(data) | |
| self.result_queue.put((idx, result)) | |
| def __init__(self, cfg, num_gpus: int = 1): | |
| """ | |
| Args: | |
| cfg (CfgNode): | |
| num_gpus (int): if 0, will run on CPU | |
| """ | |
| num_workers = max(num_gpus, 1) | |
| self.task_queue = mp.Queue(maxsize=num_workers * 3) | |
| self.result_queue = mp.Queue(maxsize=num_workers * 3) | |
| self.procs = [] | |
| for gpuid in range(max(num_gpus, 1)): | |
| cfg = cfg.clone() | |
| cfg.defrost() | |
| cfg.MODEL.DEVICE = "cuda:{}".format(gpuid) if num_gpus > 0 else "cpu" | |
| self.procs.append( | |
| AsyncPredictor._PredictWorker(cfg, self.task_queue, self.result_queue) | |
| ) | |
| self.put_idx = 0 | |
| self.get_idx = 0 | |
| self.result_rank = [] | |
| self.result_data = [] | |
| for p in self.procs: | |
| p.start() | |
| atexit.register(self.shutdown) | |
| def put(self, image): | |
| self.put_idx += 1 | |
| self.task_queue.put((self.put_idx, image)) | |
| def get(self): | |
| self.get_idx += 1 # the index needed for this request | |
| if len(self.result_rank) and self.result_rank[0] == self.get_idx: | |
| res = self.result_data[0] | |
| del self.result_data[0], self.result_rank[0] | |
| return res | |
| while True: | |
| # make sure the results are returned in the correct order | |
| idx, res = self.result_queue.get() | |
| if idx == self.get_idx: | |
| return res | |
| insert = bisect.bisect(self.result_rank, idx) | |
| self.result_rank.insert(insert, idx) | |
| self.result_data.insert(insert, res) | |
| def __len__(self): | |
| return self.put_idx - self.get_idx | |
| def __call__(self, image): | |
| self.put(image) | |
| return self.get() | |
| def shutdown(self): | |
| for _ in self.procs: | |
| self.task_queue.put(AsyncPredictor._StopToken()) | |
| def default_buffer_size(self): | |
| return len(self.procs) * 5 | |
| detrex_model_list = { | |
| # DETR | |
| "detr/detr_r50_300ep": { | |
| "configs": "projects/detr/configs/detr_r50_300ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/converted_detr_r50_500ep.pth" | |
| }, | |
| "detr/detr_r50_dc5_300ep": { | |
| "configs": "projects/detr/configs/detr_r50_dc5_300ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_detr_r50_dc5.pth" | |
| }, | |
| "detr/detr_r101_300ep.py": { | |
| "configs": "projects/detr/configs/detr_r101_300ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/converted_detr_r101_500ep.pth" | |
| }, | |
| "detr/detr_r101_dc5_300ep.py": { | |
| "configs": "projects/detr/configs/detr_r101_dc5_300ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_detr_r101_dc5.pth" | |
| }, | |
| # Deformable-DETR | |
| "deformable_detr/deformable_detr_r50_with_box_refinement_50ep": { | |
| "configs": "projects/deformable_detr/configs/deformable_detr_r50_with_box_refinement_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/deformable_detr_with_box_refinement_50ep_new.pth" | |
| }, | |
| "deformable_detr/deformable_detr_r50_two_stage_50ep": { | |
| "configs": "projects/deformable_detr/configs/deformable_detr_r50_two_stage_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/deformable_detr_r50_two_stage_50ep_new.pth" | |
| }, | |
| # Anchor-DETR | |
| "anchor_detr/anchor_detr_r50_50ep":{ | |
| "configs":"projects/anchor_detr/configs/anchor_detr_r50_50ep.py", | |
| "ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/anchor_detr_r50_50ep.pth" | |
| }, | |
| "anchor_detr/anchor_detr_r50_50ep_(converted)":{ | |
| "configs":"projects/anchor_detr/configs/anchor_detr_r50_50ep.py", | |
| "ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_anchor_detr_r50_50ep.pth" | |
| }, | |
| "anchor_detr/anchor_detr_r50_dc5_50ep":{ | |
| "configs":"projects/anchor_detr/configs/anchor_detr_r50_dc5_50ep.py", | |
| "ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_anchor_detr_r50_dc5_50ep.pth" | |
| }, | |
| "anchor_detr/anchor_detr_r101_50ep":{ | |
| "configs":"projects/anchor_detr/configs/anchor_detr_r101_dc5_50ep.py", | |
| "ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_anchor_detr_r101_dc5_50ep.pth" | |
| }, | |
| "anchor_detr/anchor_detr_r101_dc5_50ep":{ | |
| "configs":"projects/anchor_detr/configs/anchor_detr_r101_dc5_50ep.py", | |
| "ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_anchor_detr_r101_50ep.pth" | |
| }, | |
| # Conditional-DETR | |
| "conditional_detr/conditional_detr_r50_50ep":{ | |
| "configs":"projects/conditional_detr/configs/conditional_detr_r50_50ep.py", | |
| "ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/conditional_detr_r50_50ep.pth" | |
| }, | |
| "conditional_detr/conditional_detr_r50_50ep_(converted)":{ | |
| "configs":"", | |
| "ckpts":"" | |
| }, | |
| "conditional_detr/conditional_detr_r101_50ep":{ | |
| "configs":"projects/conditional_detr/configs/conditional_detr_r101_50ep.py", | |
| "ckpts":"https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/converted_conditional_detr_r101_50ep.pth" | |
| }, | |
| "conditional_detr/conditional_detr_r101_dc5_50ep": { | |
| "configs": "projects/conditional_detr/configs/conditional_detr_r101_dc5_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_conditional_detr_r101_dc5.pth" | |
| }, | |
| # DAB-DETR | |
| "dab_detr/dab_detr_r50_50ep": { | |
| "configs": "projects/dab_detr/configs/dab_detr_r50_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/dab_detr_r50_50ep.pth" | |
| }, | |
| "dab_detr/dab_detr_r50_3patterns_50ep": { | |
| "configs": "projects/dab_detr/configs/dab_detr_r50_3patterns_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_dab_detr_r50_3patterns.pth" | |
| }, | |
| "dab_detr/dab_detr_r50_dc5_50ep": { | |
| "configs": "projects/dab_detr/configs/dab_detr_r50_dc5_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_dab_detr_r50_dc5.pth" | |
| }, | |
| "dab_detr/dab_detr_r50_dc5_3patterns_50ep": { | |
| "configs": "projects/dab_detr/configs/dab_detr_r50_3patterns_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_dab_detr_r50_dc5_3patterns.pth" | |
| }, | |
| "dab_detr/dab_detr_r101_50ep": { | |
| "configs": "projects/dab_detr/configs/dab_detr_r101_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/dab_detr_r101_50ep.pth" | |
| }, | |
| "dab_detr/dab_detr_r50_dc5_3patterns_50ep_(converted)": { | |
| "configs": "projects/dab_detr/configs/dab_detr_r50_dc5_3patterns_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_detr_r101_dc5.pth" | |
| }, | |
| "dab_detr/dab_detr_swin_t_in1k_50ep": { | |
| "configs": "projects/dab_detr/configs/dab_detr_swin_t_in1k_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/dab_detr_swin_t_in1k_50ep.pth" | |
| }, | |
| "dab_detr/dab_deformable_detr_r50_50ep": { | |
| "configs": "projects/dab_detr/configs/dab_deformable_detr_r50_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dab_deformable_detr_r50_50ep_49AP.pth" | |
| }, | |
| "dab_detr/dab_deformable_detr_r50_two_stage_50ep": { | |
| "configs": "projects/dab_detr/configs/dab_deformable_detr_r50_two_stage_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dab_deformable_detr_r50_two_stage_49_7AP.pth" | |
| }, | |
| # DN-DETR | |
| "dn_detr/dn_detr_r50_50ep": { | |
| "configs": "projects/dn_detr/configs/dn_detr_r50_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.0/dn_detr_r50_50ep.pth" | |
| }, | |
| "dn_detr/dn_detr_r50_dc5_50ep": { | |
| "configs": "projects/dn_detr/configs/dn_detr_r50_dc5_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_dn_detr_r50_dc5.pth" | |
| }, | |
| # DINO | |
| "dino/dino-resnet/dino_r50_5scale_12ep": { | |
| "configs": "projects/dino/configs/dino-resnet/dino_r50_5scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_r50_5scale_12ep.pth" | |
| }, | |
| "dino/dino-resnet/dino_r50_4scale_12ep_300dn": { | |
| "configs": "projects/dino/configs/dino-resnet/dino_r50_4scale_12ep_300dn.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dino_r50_4scale_12ep_300dn.pth" | |
| }, | |
| "dino/dino-resnet/dino_r50_4scale_24ep": { | |
| "configs": "projects/dino/configs/dino-resnet/dino_r50_4scale_24ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_r50_4scale_24ep.pth" | |
| }, | |
| "dino/dino-resnet/dino_r101_4scale_12ep_": { | |
| "configs": "projects/dino/configs/dino-resnet/dino_r101_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_r101_4scale_12ep.pth" | |
| }, | |
| # Pretrained DINO with Swin-Transformer Backbone | |
| "dino/dino-swin/dino_swin_tiny_224_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_tiny_224_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_tiny_224_4scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_swin_tiny_224_22kto1k_finetune_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_tiny_224_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_tiny_224_22kto1k_finetune_4scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_swin_small_224_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_tiny_224_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_small_224_4scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_swin_base_384_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_base_384_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_base_384_4scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_swin_large_224_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_large_224_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_large_224_4scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_swin_large_384_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_large_384_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.1.1/dino_swin_large_4scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_swin_large_384_5scale_12ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_large_384_5scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_swin_large_384_5scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_swin_large_384_4scale_36ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_large_384_4scale_36ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dino_swin_large_384_4scale_36ep.pth" | |
| }, | |
| "dino/dino-swin/dino_swin_large_384_5scale_36ep": { | |
| "configs": "projects/dino/configs/dino-swin/dino_swin_large_384_5scale_36ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_swin_large_384_5scale_36ep.pth" | |
| }, | |
| # Pretrained DINO with FocalNet Backbone | |
| "dino/dino-swin/dino_focalnet_large_lrf_384_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-focal/dino_focalnet_large_lrf_384_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_focal_large_lrf_384_4scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_focalnet_large_lrf_384_fl4_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-focal/dino_focalnet_large_lrf_384_fl4_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_focal_large_lrf_384_4scale_12ep.pth" | |
| }, | |
| "dino/dino-swin/dino_focalnet_large_lrf_384_fl4_5scale_12ep": { | |
| "configs": "projects/dino/configs/dino-focal/dino_focalnet_large_lrf_384_fl4_5scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_focalnet_large_lrf_384_fl4_5scale_12ep.pth" | |
| }, | |
| # Pretrained DINO with ViTDet Backbone | |
| "dino/dino-vitdet/dino_vitdet_base_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-vitdet/dino_vitdet_base_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_vitdet_4scale_12ep.pth" | |
| }, | |
| "dino/dino-vitdet/dino_vitdet_base_4scale_50ep": { | |
| "configs": "projects/dino/configs/dino-vitdet/dino_vitdet_base_4scale_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_vitdet_base_4scale_50ep.pth" | |
| }, | |
| "dino/dino-vitdet/dino_vitdet_large_4scale_12ep": { | |
| "configs": "projects/dino/configs/dino-vitdet/dino_vitdet_large_4scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_vitdet_large_4scale_12ep.pth" | |
| }, | |
| "dino/dino-vitdet/dino_vitdet_large_4scale_50ep": { | |
| "configs": "projects/dino/configs/dino-vitdet/dino_vitdet_large_4scale_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.1/dino_vitdet_large_4scale_50ep.pth" | |
| }, | |
| # H-Deformable-DETR | |
| "h_deformable_detr/h_deformable_detr_r50_two_stage_12ep": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_r50_two_stage_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.4.0/h_deformable_detr_r50_two_stage_12ep_modified_train_net.pth" | |
| }, | |
| "h_deformable_detr/h_deformable_detr_r50_two_stage_12ep(converted)": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_r50_two_stage_12ep.py", | |
| "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" | |
| }, | |
| "h_deformable_detr/h_deformable_detr_r50_two_stage_36ep": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_r50_two_stage_36ep.py", | |
| "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" | |
| }, | |
| "h_deformable_detr/h_deformable_detr_swin_tiny_two_stage_12ep": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_tiny_two_stage_12ep.py", | |
| "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" | |
| }, | |
| "h_deformable_detr/h_deformable_detr_swin_tiny_two_stage_36ep": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_tiny_two_stage_36ep.py", | |
| "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" | |
| }, | |
| "h_deformable_detr/h_deformable_detr_swin_large_two_stage_12ep": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_large_two_stage_12ep.py", | |
| "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" | |
| }, | |
| "h_deformable_detr/h_deformable_detr_swin_large_two_stage_36ep": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_large_two_stage_36ep.py", | |
| "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" | |
| }, | |
| "h_deformable_detr/h_deformable_detr_swin_large_two_stage_12ep_900queries": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_large_two_stage_12ep_900queries.py", | |
| "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" | |
| }, | |
| "h_deformable_detr/h_deformable_detr_swin_large_two_stage_36ep_900queries": { | |
| "configs": "projects/h_deformable_detr/configs/h_deformable_detr_swin_large_two_stage_36ep_900queries.py", | |
| "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" | |
| }, | |
| # DETA | |
| "deta/improved_deformable_detr_baseline_50ep": { | |
| "configs": "projects/deta/configs/improved_deformable_detr_baseline_50ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_deta_improved_deformable_baseline.pth" | |
| }, | |
| "deta/deta_r50_5scale_12ep_bs8": { | |
| "configs": "projects/deta/configs/deta_r50_5scale_12ep_bs8.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.4.0/deta_r50_5scale_12ep_bs8.pth" | |
| }, | |
| "deta/deta_r50_5scale_12ep": { | |
| "configs": "projects/deta/configs/deta_r50_5scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.4.0/deta_r50_5scale_12ep_hacked_trainer.pth" | |
| }, | |
| "deta/deta_r50_5scale_no_frozen_backbone": { | |
| "configs": "projects/deta/configs/deta_r50_5scale_no_frozen_backbone.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.4.0/deta_r50_5scale_12ep_no_freeze_backbone.pth" | |
| }, | |
| "deta/deta_r50_5scale_12ep(converted)": { | |
| "configs": "projects/deta/configs/deta_r50_5scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_deta_r50_5scale_12ep.pth" | |
| }, | |
| "deta/DETA-Swin-Large-finetune (converted)": { | |
| "configs": "projects/deta/configs/deta_r50_5scale_12ep.py", | |
| "ckpts": "https://github.com/IDEA-Research/detrex-storage/releases/download/v0.3.0/converted_deta_swin_o365_finetune.pth" | |
| }, | |
| } | |
| def setup(args): | |
| cfg = LazyConfig.load(args.config_file) | |
| cfg = LazyConfig.apply_overrides(cfg, args.opts) | |
| return cfg | |
| def get_parser(): | |
| parser = argparse.ArgumentParser(description="detrex demo for visualizing customized inputs") | |
| parser.add_argument( | |
| "--config-file", | |
| default="projects/dino/configs/dino_r50_4scale_12ep.py", | |
| metavar="FILE", | |
| help="path to config file", | |
| ) | |
| parser.add_argument("--webcam", action="store_true", help="Take inputs from webcam.") | |
| parser.add_argument("--video-input", help="Path to video file.") | |
| parser.add_argument( | |
| "--input", | |
| nargs="+", | |
| help="A list of space separated input images; " | |
| "or a single glob pattern such as 'directory/*.jpg'", | |
| ) | |
| parser.add_argument( | |
| "--output", | |
| help="A file or directory to save output visualizations. " | |
| "If not given, will show output in an OpenCV window.", | |
| ) | |
| parser.add_argument( | |
| "--min_size_test", | |
| type=int, | |
| default=800, | |
| help="Size of the smallest side of the image during testing. Set to zero to disable resize in testing.", | |
| ) | |
| parser.add_argument( | |
| "--max_size_test", | |
| type=float, | |
| default=1333, | |
| help="Maximum size of the side of the image during testing.", | |
| ) | |
| parser.add_argument( | |
| "--img_format", | |
| type=str, | |
| default="RGB", | |
| help="The format of the loading images.", | |
| ) | |
| parser.add_argument( | |
| "--metadata_dataset", | |
| type=str, | |
| default="coco_2017_val", | |
| help="The metadata infomation to be used. Default to COCO val metadata.", | |
| ) | |
| parser.add_argument( | |
| "--confidence-threshold", | |
| type=float, | |
| default=0.5, | |
| help="Minimum score for instance predictions to be shown", | |
| ) | |
| parser.add_argument( | |
| "--opts", | |
| help="Modify config options using the command-line", | |
| default=None, | |
| nargs=argparse.REMAINDER, | |
| ) | |
| return parser | |
| 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 download_ckpts_and_image(ckpts): | |
| print("ckpts:", ckpts) | |
| torch.hub.download_url_to_file(ckpts, "dino_deitsmall16_pretrain.pth") | |
| def run_detection(input_file, output_file, model_name, input_confidence, device): | |
| configs = detrex_model_list[model_name]["configs"] | |
| ckpts = detrex_model_list[model_name]["ckpts"] | |
| mp.set_start_method("spawn", force=True) | |
| args = get_parser().parse_args([ | |
| "--config-file", configs, | |
| "--input", input_file, | |
| "--output", output_file, | |
| "--confidence-threshold", str(input_confidence), | |
| "--opts", "train.init_checkpoint=" + ckpts | |
| ]) | |
| setup_logger(name="fvcore") | |
| logger = setup_logger() | |
| logger.info("Arguments: " + str(args)) | |
| cfg = setup(args) | |
| cfg.model.device = device | |
| cfg.train.device = device | |
| model = instantiate(cfg.model) | |
| model.to(cfg.train.device) | |
| checkpointer = DetectionCheckpointer(model) | |
| checkpointer.load(cfg.train.init_checkpoint) | |
| model.eval() | |
| demo = VisualizationDemo( | |
| model=model, | |
| min_size_test=args.min_size_test, | |
| max_size_test=args.max_size_test, | |
| img_format=args.img_format, | |
| metadata_dataset=args.metadata_dataset, | |
| ) | |
| if args.input: | |
| if len(args.input) == 1: | |
| args.input = glob.glob(os.path.expanduser(args.input[0])) | |
| assert args.input, "The input path(s) was not found" | |
| for path in tqdm.tqdm(args.input, disable=not args.output): | |
| # use PIL, to be consistent with evaluation | |
| img = read_image(path, format="BGR") | |
| start_time = time.time() | |
| predictions, visualized_output = demo.run_on_image(img, args.confidence_threshold) | |
| logger.info( | |
| "{}: {} in {:.2f}s".format( | |
| path, | |
| "detected {} instances".format(len(predictions["instances"])) | |
| if "instances" in predictions | |
| else "finished", | |
| time.time() - start_time, | |
| ) | |
| ) | |
| if args.output: | |
| if os.path.isdir(args.output): | |
| assert os.path.isdir(args.output), args.output | |
| out_filename = os.path.join(args.output, os.path.basename(path)) | |
| else: | |
| assert len(args.input) == 1, "Please specify a directory with args.output" | |
| out_filename = args.output | |
| visualized_output.save(out_filename) | |
| def download_test_img(): | |
| import shutil | |
| torch.hub.download_url_to_file( | |
| 'https://github.com/isLinXu/issues/files/12658779/projects.zip', | |
| 'projects.zip') | |
| # Images | |
| torch.hub.download_url_to_file( | |
| 'https://user-images.githubusercontent.com/59380685/268517006-d8d4d3b3-964a-4f4d-8458-18c7eb75a4f2.jpg', | |
| '000000502136.jpg') | |
| shutil.unpack_archive('projects.zip', './', 'zip') | |
| def detect_image(input_image, model_name, input_confidence, device): | |
| input_dir = "input.jpg" | |
| input_image.save(input_dir) | |
| output_image = "output.jpg" | |
| run_detection(input_dir, output_image, model_name, input_confidence, device) | |
| return output_image | |
| if __name__ == '__main__': | |
| input_image = gr.inputs.Image(type='pil', label="Input Image") | |
| input_model_name = gr.inputs.Dropdown(list(detrex_model_list.keys()), label="Model Name", default="dab_detr/dab_detr_r50_50ep") | |
| input_confidence = gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.25, label="Confidence Threshold") | |
| input_device = gr.inputs.Radio(["cpu", "cuda"], label="Device", default="cpu") | |
| output_image = gr.outputs.Image(type='pil', label="Output Image") | |
| download_test_img() | |
| examples = [["000000502136.jpg", "dab_detr/dab_detr_r50_50ep", 0.25, "cpu"]] | |
| title = "🦖detrex: Benchmarking Detection Transformers web demo" | |
| description = "<div align='center'><img src='https://raw.githubusercontent.com/IDEA-Research/detrex/main/assets/logo_2.png' width='250''/><div>" \ | |
| "<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）。" \ | |
| "detrex is a research platform for DETR-based object detection, segmentation, pose estimation and other visual recognition tasks.</p>" | |
| article = "<p style='text-align: center'><a href='https://github.com/IDEA-Research/detrex'>detrex</a></p>" \ | |
| "<p style='text-align: center'><a href='https://github.com/isLinXu'>gradio build by gatilin</a></a></p>" | |
| image_interface = gr.Interface(detect_image, | |
| inputs=[input_image, input_model_name, input_confidence, input_device], | |
| outputs=output_image,examples=examples, | |
| title=title, article=article, description=description) | |
| image_interface.launch() | |