# Adapted from https://huggingface.co/spaces/hlydecker/MegaDetector_v5 # Adapted from https://huggingface.co/spaces/sofmi/MegaDetector_DLClive/blob/main/app.py # Adapted from https://huggingface.co/spaces/Neslihan/megadetector_dlcmodels/blob/main/app.py # Adapted from https://huggingface.co/spaces/DeepLabCut/MegaDetector_DeepLabCut import os import yaml import numpy as np from matplotlib import cm import gradio as gr import deeplabcut import dlclibrary import dlclive import transformers from PIL import Image, ImageColor, ImageFont, ImageDraw import requests from viz_utils import save_results_as_json, draw_keypoints_on_image, draw_bbox_w_text, save_results_only_dlc from detection_utils import predict_md, crop_animal_detections from dlc_utils import predict_dlc from ui_utils import gradio_inputs_for_MD_DLC, gradio_outputs_for_MD_DLC, gradio_description_and_examples from deeplabcut.utils import auxiliaryfunctions from dlclibrary.dlcmodelzoo.modelzoo_download import ( download_huggingface_model, MODELOPTIONS, ) from dlclive import DLCLive, Processor # TESTING (passes) download the SuperAnimal models: #model = 'superanimal_topviewmouse' #train_dir = 'DLC_models/sa-tvm' #download_huggingface_model(model, train_dir) # grab demo data cooco cat: url = "http://images.cocodataset.org/val2017/000000039769.jpg" image = Image.open(requests.get(url, stream=True).raw) # megadetector and dlc model look up MD_models_dict = {'md_v5a': "MD_models/md_v5a.0.0.pt", # 'md_v5b': "MD_models/md_v5b.0.0.pt"} # DLC models target dirs DLC_models_dict = {'superanimal_topviewmouse_dlcrnet': "DLC_models/sa-tvm", 'superanimal_quadruped_dlcrnet': "DLC_models/sa-q"} ##################################################### def predict_pipeline(img_input, mega_model_input, dlc_model_input_str, flag_dlc_only, flag_show_str_labels, bbox_likelihood_th, kpts_likelihood_th, font_style, font_size, keypt_color, marker_size, ): if not flag_dlc_only: ############################################################ # ### Run Megadetector md_results = predict_md(img_input, MD_models_dict[mega_model_input], #mega_model_input, size=640) #Image.fromarray(results.imgs[0]) ################################################################ # Obtain animal crops for bboxes with confidence above th list_crops = crop_animal_detections(img_input, md_results, bbox_likelihood_th) ############################################################ ## Get DLC model and label map # If model is found: do not download (previous execution is likely within same day) # TODO: can we ask the user whether to reload dlc model if a directory is found? if os.path.isdir(DLC_models_dict[dlc_model_input_str]) and \ len(os.listdir(DLC_models_dict[dlc_model_input_str])) > 0: path_to_DLCmodel = DLC_models_dict[dlc_model_input_str] else: path_to_DLCmodel = DLC_models_dict[dlc_model_input_str] download_huggingface_model(dlc_model_input_str, path_to_DLCmodel) # extract map label ids to strings pose_cfg_path = os.path.join(DLC_models_dict[dlc_model_input_str], 'pose_cfg.yaml') with open(pose_cfg_path, "r") as stream: pose_cfg_dict = yaml.safe_load(stream) map_label_id_to_str = dict([(k,v) for k,v in zip([el[0] for el in pose_cfg_dict['all_joints']], # pose_cfg_dict['all_joints'] is a list of one-element lists, pose_cfg_dict['all_joints_names'])]) ############################################################## # Run DLC and visualize results dlc_proc = Processor() #TODO: update deeplabcut.video_inference_superanimal() once merged # if required: ignore MD crops and run DLC on full image [mostly for testing] if flag_dlc_only: # compute kpts on input img list_kpts_per_crop = predict_dlc([np.asarray(img_input)], kpts_likelihood_th, path_to_DLCmodel, dlc_proc) # draw kpts on input img #fix! draw_keypoints_on_image(img_input, list_kpts_per_crop[0], # a numpy array with shape [num_keypoints, 2]. map_label_id_to_str, flag_show_str_labels, use_normalized_coordinates=False, font_style=font_style, font_size=font_size, keypt_color=keypt_color, marker_size=marker_size) donw_file = save_results_only_dlc(list_kpts_per_crop[0], map_label_id_to_str,dlc_model_input_str) return img_input, donw_file else: # Compute kpts for each crop list_kpts_per_crop = predict_dlc(list_crops, kpts_likelihood_th, path_to_DLCmodel, dlc_proc) # resize input image to match megadetector output img_background = img_input.resize((md_results.ims[0].shape[1], md_results.ims[0].shape[0])) # draw keypoints on each crop and paste to background img for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops, list_kpts_per_crop)): img_crop = Image.fromarray(np_crop) # Draw keypts on crop draw_keypoints_on_image(img_crop, kpts_crop, # a numpy array with shape [num_keypoints, 2]. map_label_id_to_str, flag_show_str_labels, use_normalized_coordinates=False, # if True, then I should use md_results.xyxyn for list_kpts_crop font_style=font_style, font_size=font_size, keypt_color=keypt_color, marker_size=marker_size) # Paste crop in original image img_background.paste(img_crop, box = tuple([int(t) for t in md_results.xyxy[0][ic,:2]])) # Plot bbox bb_per_animal = md_results.xyxy[0].tolist()[ic] pred = md_results.xyxy[0].tolist()[ic][4] if bbox_likelihood_th < pred: draw_bbox_w_text(img_background, bb_per_animal, font_size=font_size) # TODO: add selectable color for bbox? # Save detection results as json download_file = save_results_as_json(md_results,list_kpts_per_crop,map_label_id_to_str, bbox_likelihood_th,dlc_model_input_str,mega_model_input) return img_background, download_file ######################################################### # Define user interface and launch inputs = gradio_inputs_for_MD_DLC(list(MD_models_dict.keys()), list(DLC_models_dict.keys())) outputs = gradio_outputs_for_MD_DLC() [gr_title, gr_description, examples] = gradio_description_and_examples() # launch demo = gr.Interface(predict_pipeline, inputs=inputs, outputs=outputs, title=gr_title, description=gr_description, examples=examples, theme="huggingface") demo.launch(enable_queue=True, share=True)