Spaces:
Running on T4
Running on T4
| import os | |
| from typing import TypeVar | |
| from tqdm import tqdm | |
| import gradio as gr | |
| import numpy as np | |
| import supervision as sv | |
| from PIL import Image | |
| from rfdetr import ( | |
| RFDETRNano, RFDETRSmall, RFDETRMedium, RFDETRBase, RFDETRLarge, | |
| RFDETRSegNano, RFDETRSegSmall, RFDETRSegMedium, | |
| RFDETRSegLarge, RFDETRSegXLarge, RFDETRSeg2XLarge, | |
| ) | |
| from rfdetr.detr import RFDETR | |
| from rfdetr.util.coco_classes import COCO_CLASSES | |
| from utils.image import calculate_resolution_wh | |
| from utils.video import create_directory, generate_unique_name | |
| ImageType = TypeVar("ImageType", Image.Image, np.ndarray) | |
| MARKDOWN = """ | |
| # RF-DETR 🔥 | |
| [`[code]`](https://github.com/roboflow/rf-detr) | |
| [`[blog]`](https://blog.roboflow.com/rf-detr) | |
| [`[notebook]`](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-finetune-rf-detr-on-detection-dataset.ipynb) | |
| RF-DETR is a real-time, transformer-based object detection model architecture developed | |
| by [Roboflow](https://roboflow.com/) and released under the Apache 2.0 license. | |
| """ | |
| IMAGE_PROCESSING_EXAMPLES = [ | |
| ['https://media.roboflow.com/supervision/image-examples/people-walking.png', 0.3, 1024, "medium (object detection)"], | |
| ['https://media.roboflow.com/supervision/image-examples/vehicles.png', 0.3, 1024, "medium (object detection)"], | |
| ['https://media.roboflow.com/supervision/image-examples/motorbike.png', 0.3, 1024, "medium (object detection)"], | |
| ['https://media.roboflow.com/notebooks/examples/dog-2.jpeg', 0.5, 512, "nano (object detection)"], | |
| ['https://media.roboflow.com/notebooks/examples/dog-3.jpeg', 0.5, 512, "nano (object detection)"], | |
| ['https://media.roboflow.com/supervision/image-examples/basketball-1.png', 0.5, 512, "nano (object detection)"], | |
| ['https://media.roboflow.com/notebooks/examples/dog-2.jpeg', 0.5, 512, "medium (instance segmentation)"], | |
| ['https://media.roboflow.com/notebooks/examples/dog-3.jpeg', 0.5, 512, "medium (instance segmentation)"], | |
| ['https://media.roboflow.com/supervision/image-examples/basketball-1.png', 0.5, 512, "medium (instance segmentation)"], | |
| ] | |
| VIDEO_PROCESSING_EXAMPLES = [ | |
| ["videos/people-walking.mp4", 0.3, 1024, "medium (object detection)"], | |
| ["videos/vehicles.mp4", 0.3, 1024, "medium (object detection)"], | |
| ] | |
| COLOR = sv.ColorPalette.from_hex([ | |
| "#ffff00", "#ff9b00", "#ff8080", "#ff66b2", "#ff66ff", "#b266ff", | |
| "#9999ff", "#3399ff", "#66ffff", "#33ff99", "#66ff66", "#99ff00" | |
| ]) | |
| MAX_VIDEO_LENGTH_SECONDS = 5 | |
| VIDEO_SCALE_FACTOR = 0.5 | |
| VIDEO_TARGET_DIRECTORY = "tmp" | |
| create_directory(directory_path=VIDEO_TARGET_DIRECTORY) | |
| def detect_and_annotate( | |
| model: RFDETR, | |
| image: ImageType, | |
| confidence: float, | |
| checkpoint: str = "medium (object detection)" | |
| ) -> ImageType: | |
| detections = model.predict(image, threshold=confidence) | |
| resolution_wh = calculate_resolution_wh(image) | |
| text_scale = sv.calculate_optimal_text_scale(resolution_wh=resolution_wh) - 0.2 | |
| thickness = sv.calculate_optimal_line_thickness(resolution_wh=resolution_wh) | |
| mask_annotator = sv.MaskAnnotator(color=COLOR) | |
| bbox_annotator = sv.BoxAnnotator(color=COLOR, thickness=thickness) | |
| label_annotator = sv.LabelAnnotator( | |
| color=COLOR, | |
| text_color=sv.Color.BLACK, | |
| text_scale=text_scale | |
| ) | |
| labels = [ | |
| f"{COCO_CLASSES[class_id]} {confidence:.2f}" | |
| for class_id, confidence | |
| in zip(detections.class_id, detections.confidence) | |
| ] | |
| print(detections) | |
| annotated_image = image.copy() | |
| if checkpoint in SEGMENTATION_CHECKPOINTS: | |
| annotated_image = mask_annotator.annotate(annotated_image, detections) | |
| annotated_image = bbox_annotator.annotate(annotated_image, detections) | |
| annotated_image = label_annotator.annotate(annotated_image, detections, labels) | |
| return annotated_image | |
| def load_model(resolution: int, checkpoint: str) -> RFDETR: | |
| if checkpoint == "nano (object detection)": | |
| return RFDETRNano(resolution=resolution) | |
| if checkpoint == "small (object detection)": | |
| return RFDETRSmall(resolution=resolution) | |
| if checkpoint == "medium (object detection)": | |
| return RFDETRMedium(resolution=resolution) | |
| if checkpoint == "base (object detection)": | |
| return RFDETRBase(resolution=resolution) | |
| if checkpoint == "large (object detection)": | |
| return RFDETRLarge(resolution=resolution) | |
| if checkpoint == "nano (instance segmentation)": | |
| return RFDETRSegNano(resolution=resolution) | |
| if checkpoint == "small (instance segmentation)": | |
| return RFDETRSegSmall(resolution=resolution) | |
| if checkpoint == "medium (instance segmentation)": | |
| return RFDETRSegMedium(resolution=resolution) | |
| if checkpoint == "large (instance segmentation)": | |
| return RFDETRSegLarge(resolution=resolution) | |
| if checkpoint == "xlarge (instance segmentation)": | |
| return RFDETRSegXLarge(resolution=resolution) | |
| if checkpoint == "2xlarge (instance segmentation)": | |
| return RFDETRSeg2XLarge(resolution=resolution) | |
| raise TypeError(f"Unknown checkpoint: {checkpoint}") | |
| SEGMENTATION_CHECKPOINTS = { | |
| "nano (instance segmentation)", | |
| "small (instance segmentation)", | |
| "medium (instance segmentation)", | |
| "large (instance segmentation)", | |
| "xlarge (instance segmentation)", | |
| "2xlarge (instance segmentation)", | |
| } | |
| def adjust_resolution(checkpoint: str, resolution: int) -> int: | |
| if checkpoint in SEGMENTATION_CHECKPOINTS: | |
| divisor = 24 | |
| elif checkpoint in {"nano (object detection)", "small (object detection)", "medium (object detection)"}: | |
| divisor = 32 | |
| elif checkpoint in {"base (object detection)", "large (object detection)"}: | |
| divisor = 56 | |
| else: | |
| raise ValueError(f"Unknown checkpoint: {checkpoint}") | |
| remainder = resolution % divisor | |
| if remainder == 0: | |
| return resolution | |
| lower = resolution - remainder | |
| upper = lower + divisor | |
| if resolution - lower < upper - resolution: | |
| return lower | |
| else: | |
| return upper | |
| def image_processing_inference( | |
| input_image: Image.Image, | |
| confidence: float, | |
| resolution: int, | |
| checkpoint: str | |
| ): | |
| resolution = adjust_resolution(checkpoint=checkpoint, resolution=resolution) | |
| model = load_model(resolution=resolution, checkpoint=checkpoint) | |
| return detect_and_annotate(model=model, image=input_image, confidence=confidence, checkpoint=checkpoint) | |
| def video_processing_inference( | |
| input_video: str, | |
| confidence: float, | |
| resolution: int, | |
| checkpoint: str, | |
| ): | |
| resolution = adjust_resolution(checkpoint=checkpoint, resolution=resolution) | |
| model = load_model(resolution=resolution, checkpoint=checkpoint) | |
| name = generate_unique_name() | |
| output_video = os.path.join(VIDEO_TARGET_DIRECTORY, f"{name}.mp4") | |
| video_info = sv.VideoInfo.from_video_path(input_video) | |
| video_info.width = int(video_info.width * VIDEO_SCALE_FACTOR) | |
| video_info.height = int(video_info.height * VIDEO_SCALE_FACTOR) | |
| total = min(video_info.total_frames, video_info.fps * MAX_VIDEO_LENGTH_SECONDS) | |
| frames_generator = sv.get_video_frames_generator(input_video, end=total) | |
| with sv.VideoSink(output_video, video_info=video_info) as sink: | |
| for frame in tqdm(frames_generator, total=total): | |
| annotated_frame = detect_and_annotate( | |
| model=model, | |
| image=frame, | |
| confidence=confidence, | |
| checkpoint=checkpoint | |
| ) | |
| annotated_frame = sv.scale_image(annotated_frame, VIDEO_SCALE_FACTOR) | |
| sink.write_frame(annotated_frame) | |
| return output_video | |
| with gr.Blocks() as demo: | |
| gr.Markdown(MARKDOWN) | |
| with gr.Tab("Image"): | |
| with gr.Row(): | |
| image_processing_input_image = gr.Image( | |
| label="Upload image", | |
| image_mode='RGB', | |
| type='pil', | |
| height=600 | |
| ) | |
| image_processing_output_image = gr.Image( | |
| label="Output image", | |
| image_mode='RGB', | |
| type='pil', | |
| height=600 | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_processing_confidence_slider = gr.Slider( | |
| label="Confidence", | |
| minimum=0.0, | |
| maximum=1.0, | |
| step=0.05, | |
| value=0.5, | |
| ) | |
| image_processing_resolution_slider = gr.Slider( | |
| label="Inference resolution", | |
| minimum=224, | |
| maximum=2240, | |
| step=1, | |
| value=896, | |
| ) | |
| image_processing_checkpoint_dropdown = gr.Dropdown( | |
| label="Checkpoint", | |
| choices=[ | |
| "nano (object detection)", | |
| "small (object detection)", | |
| "medium (object detection)", | |
| "nano (instance segmentation)", | |
| "small (instance segmentation)", | |
| "medium (instance segmentation)", | |
| "large (instance segmentation)", | |
| "xlarge (instance segmentation)", | |
| "2xlarge (instance segmentation)", | |
| ], | |
| value="medium (object detection)" | |
| ) | |
| with gr.Column(): | |
| image_processing_submit_button = gr.Button("Submit", value="primary") | |
| gr.Examples( | |
| fn=image_processing_inference, | |
| examples=IMAGE_PROCESSING_EXAMPLES, | |
| inputs=[ | |
| image_processing_input_image, | |
| image_processing_confidence_slider, | |
| image_processing_resolution_slider, | |
| image_processing_checkpoint_dropdown | |
| ], | |
| outputs=image_processing_output_image, | |
| ) | |
| image_processing_submit_button.click( | |
| image_processing_inference, | |
| inputs=[ | |
| image_processing_input_image, | |
| image_processing_confidence_slider, | |
| image_processing_resolution_slider, | |
| image_processing_checkpoint_dropdown | |
| ], | |
| outputs=image_processing_output_image, | |
| ) | |
| with gr.Tab("Video"): | |
| with gr.Row(): | |
| video_processing_input_video = gr.Video( | |
| label='Upload video', | |
| height=600 | |
| ) | |
| video_processing_output_video = gr.Video( | |
| label='Output video', | |
| height=600 | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| video_processing_confidence_slider = gr.Slider( | |
| label="Confidence", | |
| minimum=0.0, | |
| maximum=1.0, | |
| step=0.05, | |
| value=0.5, | |
| ) | |
| video_processing_resolution_slider = gr.Slider( | |
| label="Inference resolution", | |
| minimum=560, | |
| maximum=1120, | |
| step=56, | |
| value=728, | |
| ) | |
| video_processing_checkpoint_dropdown = gr.Dropdown( | |
| label="Checkpoint", | |
| choices=[ | |
| "nano (object detection)", | |
| "small (object detection)", | |
| "medium (object detection)", | |
| "nano (instance segmentation)", | |
| "small (instance segmentation)", | |
| "medium (instance segmentation)", | |
| "large (instance segmentation)", | |
| "xlarge (instance segmentation)", | |
| "2xlarge (instance segmentation)", | |
| ], | |
| value="medium (object detection)" | |
| ) | |
| with gr.Column(): | |
| video_processing_submit_button = gr.Button("Submit", value="primary") | |
| gr.Examples( | |
| fn=video_processing_inference, | |
| examples=VIDEO_PROCESSING_EXAMPLES, | |
| inputs=[ | |
| video_processing_input_video, | |
| video_processing_confidence_slider, | |
| video_processing_resolution_slider, | |
| video_processing_checkpoint_dropdown | |
| ], | |
| outputs=video_processing_output_video | |
| ) | |
| video_processing_submit_button.click( | |
| video_processing_inference, | |
| inputs=[ | |
| video_processing_input_video, | |
| video_processing_confidence_slider, | |
| video_processing_resolution_slider, | |
| video_processing_checkpoint_dropdown | |
| ], | |
| outputs=video_processing_output_video | |
| ) | |
| demo.launch(debug=False, show_error=True) | |