import gradio as gr from matplotlib import colormaps from matplotlib.colors import to_hex from PIL import Image from pytorch_utils import MAX_IMAGE_SIZE from viz_utils import COLORMAPS # shades built around the DeepLabCut docs palette (DeepLabCut/docs/_static/custom.css) DLC_PURPLE = gr.themes.Color( name="dlc_purple", c50="#f5edff", c100="#ead7ff", c200="#d9b8ff", c300="#c084fc", c400="#ac72f0", c500="#9b5de5", c600="#8550c4", c700="#73439a", c800="#5c357c", c900="#4b236f", c950="#2e1546", ) DLC_TEAL = gr.themes.Color( name="dlc_teal", c50="#e8f8f7", c100="#c9efec", c200="#9fe2dd", c300="#7ad3ce", c400="#57c4be", c500="#2fb0a8", c600="#21a197", c700="#1a8078", c800="#16655f", c900="#124f4a", c950="#0a2e2b", ) def dlc_theme(): # white text on purple 500 is 4.1:1, so filled buttons use 700 (7.0:1) and 600 (5.3:1) return gr.themes.Default(primary_hue=DLC_PURPLE, secondary_hue=DLC_TEAL, neutral_hue="slate").set( button_primary_background_fill="*primary_700", button_primary_background_fill_hover="*primary_600", button_primary_background_fill_dark="*primary_600", button_primary_background_fill_hover_dark="*primary_500", button_primary_text_color="white", button_primary_text_color_dark="white", button_primary_border_color="*primary_700", button_primary_border_color_dark="*primary_600", ) def gradio_inputs_for_MD_DLC(backends_list, md_models_list, dlc_models_list): # Input image gr_image_input = gr.Image(type="pil", label="Input Image") # Models gr_backend_input = gr.Radio( choices=backends_list, value=backends_list[0], label="Select backend", ) gr_mega_model_input = gr.Dropdown( choices=md_models_list, value="md_v5a", type="value", label="Select Detector model (TensorFlow legacy only)", visible=gr_backend_input.value != "PyTorch", ) gr_dlc_model_input = gr.Dropdown( choices=dlc_models_list, value="superanimal_quadruped", type="value", label="Select DeepLabCut model", ) # Other inputs gr_dlc_only_checkbox = gr.Checkbox( value=False, label="Run DeepLabCut only, directly on input image?", ) # Gradio Slider signature is (minimum, maximum, value, step, ...) gr_slider_conf_bboxes = gr.Slider( minimum=0, maximum=1, value=0.2, step=0.05, label="Set confidence threshold for animal detections", ) gr_slider_conf_keypoints = gr.Slider( minimum=0, maximum=1, value=0.4, step=0.05, label="Set confidence threshold for keypoints", ) # Data viz with gr.Accordion("Display options", open=False): gr_str_labels_checkbox = gr.Checkbox( value=True, label="Show bodypart labels?", ) gr_color_by_confidence_checkbox = gr.Checkbox( value=False, label="Color keypoints by confidence? (otherwise by bodypart)", ) gr_colormap = gr.Dropdown( choices=COLORMAPS, value="viridis", type="value", label="Keypoint colormap", ) gr_keypt_color = gr.ColorPicker( value="#862db7", label="Choose color for keypoint label", ) gr_bbox_color = gr.ColorPicker( value="#ff0000", label="Choose color for bounding boxes", ) gr_labels_font_style = gr.Dropdown( choices=["amiko", "animals", "nature", "painter", "zen"], value="amiko", type="value", label="Select keypoint label font", ) gr_slider_font_size = gr.Slider( minimum=5, maximum=30, value=18, step=1, label="Set font size", ) gr_slider_marker_size = gr.Slider( minimum=1, maximum=20, value=6, step=1, label="Set marker size", ) return [ gr_image_input, gr_backend_input, gr_mega_model_input, gr_dlc_model_input, gr_dlc_only_checkbox, gr_str_labels_checkbox, gr_slider_conf_bboxes, gr_slider_conf_keypoints, gr_labels_font_style, gr_slider_font_size, gr_keypt_color, gr_slider_marker_size, gr_color_by_confidence_checkbox, gr_colormap, gr_bbox_color, ] def confidence_legend_html(colormap="viridis"): # the colormap as a CSS gradient, matching the keypoint fill in draw_keypoints_on_image gradient = ", ".join(f"{to_hex(colormaps[colormap](i / 10))} {i * 10}%" for i in range(11)) # no leading newline: gradio prefixes "'" to cached example values starting with one (CSV injection guard) return f"""
Keypoint confidence
00.51
""" def gradio_outputs_for_MD_DLC(): gr_image_output = gr.Image(type="pil", label="Output Image") gr_confidence_legend = gr.HTML("") with gr.Row(): gr_file_download = gr.File(label="Download JSON file") gr_image_download = gr.File(label="Download annotated image") gr_confidence_table = gr.Dataframe( headers=["animal", "bodypart", "confidence"], label="Keypoint confidence (lowest first)", interactive=False, ) return [gr_image_output, gr_confidence_legend, gr_file_download, gr_image_download, gr_confidence_table] def example_sizes(path): # font and marker sizes for the resolution the PyTorch backend draws on side = min(max(Image.open(path).size), MAX_IMAGE_SIZE) return round(side / 64), round(side / 200) def gradio_description_and_examples(): title = "DeepLabCut Model Zoo: SuperAnimals" description = ( "Estimate animal poses with the SuperAnimal models from the " "[DeepLabCut Model Zoo](http://www.mackenziemathislab.org/dlc-modelzoo) " "([paper](https://arxiv.org/abs/2203.07436)). " "Upload an image or pick an example below; to run on videos, see the Model Zoo page." ) examples = [ [image, "PyTorch", "md_v5a", "superanimal_quadruped", False, True, 0.5, 0.4, "amiko"] + [font_size, "#ff0000", marker_size, False, "viridis", "#ff0000"] for image in ( "examples/dog.jpeg", "examples/cat.jpg", ) for font_size, marker_size in [example_sizes(image)] ] return [title, description, examples]