# Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. """ Gradio Demo for image detection""" # Importing necessary basic libraries and modules import os # PyTorch imports import torch # Importing the model, dataset, transformations and utility functions from PytorchWildlife from PytorchWildlife.models import detection as pw_detection from PytorchWildlife import utils as pw_utils # Importing basic libraries import shutil import time from PIL import Image import supervision as sv import gradio as gr from zipfile import ZipFile from torch.utils.data import DataLoader import numpy as np import ast # Importing the models, dataset, transformations, and utility functions from PytorchWildlife from PytorchWildlife.models import classification as pw_classification from PytorchWildlife.data import transforms as pw_trans from PytorchWildlife.data import datasets as pw_data # Setting the device to use for computations ('cuda' indicates GPU) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Initializing a supervision box annotator for visualizing detections dot_annotator = sv.DotAnnotator(radius=6) box_annotator = sv.BoxAnnotator(thickness=4) lab_annotator = sv.LabelAnnotator(text_color=sv.Color.BLACK, text_thickness=4, text_scale=2) # Create a temp folder os.makedirs(os.path.join("..","temp"), exist_ok=True) # ASK: Why do we need this? # Initializing the detection and classification models detection_model = None classification_model = None # Defining functions for different detection scenarios def load_models(det, version, clf, wpath=None, wclass=None): global detection_model, classification_model if det != "None": if det == "HerdNet General": detection_model = pw_detection.HerdNet(device=DEVICE) elif det == "HerdNet Ennedi": detection_model = pw_detection.HerdNet(device=DEVICE, version="ennedi") else: detection_model = pw_detection.__dict__[det](device=DEVICE, pretrained=True, version=version) else: detection_model = None return "NO MODEL LOADED!!" if clf != "None": # Create an exception for custom weights if clf == "CustomWeights": if (wpath is not None) and (wclass is not None): wclass = ast.literal_eval(wclass) classification_model = pw_classification.__dict__[clf](weights=wpath, class_names=wclass, device=DEVICE) else: classification_model = pw_classification.__dict__[clf](device=DEVICE, pretrained=True) else: classification_model = None return "Loaded Detector: {}. Version: {}. Loaded Classifier: {}".format(det, version, clf) def single_image_detection(input_img, det_conf_thres, clf_conf_thres, img_index=None): """Performs detection on a single image and returns an annotated image. Args: input_img (PIL.Image): Input image in PIL.Image format defaulted by Gradio. det_conf_thres (float): Confidence threshold for detection. clf_conf_thres (float): Confidence threshold for classification. img_index: Image index identifier. Returns: annotated_img (PIL.Image.Image): Annotated image with bounding box instances. """ input_img = np.array(input_img) # If the detection model is HerdNet, use dot annotator, else use box annotator if detection_model.__class__.__name__.__contains__("HerdNet"): annotator = dot_annotator # Herdnet receives both clf and det confidence thresholds results_det = detection_model.single_image_detection(input_img, img_path=img_index, det_conf_thres=det_conf_thres, clf_conf_thres=clf_conf_thres) else: annotator = box_annotator results_det = detection_model.single_image_detection(input_img, img_path=img_index, det_conf_thres = det_conf_thres) if classification_model is not None: labels = [] for xyxy, det_id in zip(results_det["detections"].xyxy, results_det["detections"].class_id): # Only run classifier when detection class is animal if det_id == 0: cropped_image = sv.crop_image(image=input_img, xyxy=xyxy) results_clf = classification_model.single_image_classification(cropped_image) labels.append("{} {:.2f}".format(results_clf["prediction"] if results_clf["confidence"] > clf_conf_thres else "Unknown", results_clf["confidence"])) else: labels = results_det["labels"] else: labels = results_det["labels"] annotated_img = lab_annotator.annotate( scene=annotator.annotate( scene=input_img, detections=results_det["detections"], ), detections=results_det["detections"], labels=labels, ) return annotated_img def batch_detection(zip_file, timelapse, det_conf_thres): """Perform detection on a batch of images from a zip file and return path to results JSON. Args: zip_file (File): Zip file containing images. det_conf_thres (float): Confidence threshold for detection. timelapse (boolean): Flag to output JSON for timelapse. clf_conf_thres (float): Confidence threshold for classification. Returns: json_save_path (str): Path to the JSON file containing detection results. """ # Clean the temp folder if it contains files extract_path = os.path.join("..","temp","zip_upload") if os.path.exists(extract_path): shutil.rmtree(extract_path) os.makedirs(extract_path) json_save_path = os.path.join(extract_path, "results.json") with ZipFile(zip_file.name) as zfile: zfile.extractall(extract_path) # Check the contents of the extracted folder extracted_files = os.listdir(extract_path) if len(extracted_files) == 1 and os.path.isdir(os.path.join(extract_path, extracted_files[0])): tgt_folder_path = os.path.join(extract_path, extracted_files[0]) else: tgt_folder_path = extract_path # If the detection model is HerdNet set batch_size to 1 if detection_model.__class__.__name__.__contains__("HerdNet"): det_results = detection_model.batch_image_detection(tgt_folder_path, batch_size=1, det_conf_thres=det_conf_thres, id_strip=tgt_folder_path) else: det_results = detection_model.batch_image_detection(tgt_folder_path, batch_size=16, det_conf_thres=det_conf_thres, id_strip=tgt_folder_path) if classification_model is not None: clf_dataset = pw_data.DetectionCrops( det_results, transform=pw_trans.Classification_Inference_Transform(target_size=224), path_head=tgt_folder_path ) clf_loader = DataLoader(clf_dataset, batch_size=32, shuffle=False, pin_memory=True, num_workers=4, drop_last=False) clf_results = classification_model.batch_image_classification(clf_loader, id_strip=tgt_folder_path) if timelapse: json_save_path = json_save_path.replace(".json", "_timelapse.json") pw_utils.save_detection_classification_timelapse_json(det_results=det_results, clf_results=clf_results, det_categories=detection_model.CLASS_NAMES, clf_categories=classification_model.CLASS_NAMES, output_path=json_save_path) else: pw_utils.save_detection_classification_json(det_results=det_results, clf_results=clf_results, det_categories=detection_model.CLASS_NAMES, clf_categories=classification_model.CLASS_NAMES, output_path=json_save_path) else: if timelapse: json_save_path = json_save_path.replace(".json", "_timelapse.json") pw_utils.save_detection_timelapse_json(det_results, json_save_path, categories=detection_model.CLASS_NAMES) elif detection_model.__class__.__name__.__contains__("HerdNet"): pw_utils.save_detection_json_as_dots(det_results, json_save_path, categories=detection_model.CLASS_NAMES) else: pw_utils.save_detection_json(det_results, json_save_path, categories=detection_model.CLASS_NAMES) return json_save_path def batch_path_detection(tgt_folder_path, det_conf_thres): """Perform detection on a batch of images from a zip file and return path to results JSON. Args: tgt_folder_path (str): path to the folder containing the images. det_conf_thres (float): Confidence threshold for detection. Returns: json_save_path (str): Path to the JSON file containing detection results. """ json_save_path = os.path.join(tgt_folder_path, "results.json") det_results = detection_model.batch_image_detection(tgt_folder_path, det_conf_thres=det_conf_thres, id_strip=tgt_folder_path) if detection_model.__class__.__name__.__contains__("HerdNet"): pw_utils.save_detection_json_as_dots(det_results, json_save_path, categories=detection_model.CLASS_NAMES) else: pw_utils.save_detection_json(det_results, json_save_path, categories=detection_model.CLASS_NAMES) return json_save_path def video_detection(video, det_conf_thres, clf_conf_thres, target_fps, codec): """Perform detection on a video and return path to processed video. Args: video (str): Video source path. det_conf_thres (float): Confidence threshold for detection. clf_conf_thres (float): Confidence threshold for classification. """ def callback(frame, index): annotated_frame = single_image_detection(frame, img_index=index, det_conf_thres=det_conf_thres, clf_conf_thres=clf_conf_thres) return annotated_frame target_path = os.path.join("..","temp","video_detection.mp4") pw_utils.process_video(source_path=video, target_path=target_path, callback=callback, target_fps=int(target_fps), codec=codec) return target_path # Building Gradio UI with gr.Blocks() as demo: gr.Markdown("# Pytorch-Wildlife Demo.") with gr.Row(): det_drop = gr.Dropdown( ["None", "MegaDetectorV5", "MegaDetectorV6", "HerdNet General", "HerdNet Ennedi"], label="Detection model", info="Will add more detection models!", value="None" # Default ) det_version = gr.Dropdown( ["None"], label="Model version", info="Select the version of the model", value="None", ) with gr.Column(): clf_drop = gr.Dropdown( ["None", "AI4GOpossum", "AI4GAmazonRainforest", "AI4GSnapshotSerengeti", "CustomWeights"], interactive=True, label="Classification model", info="Will add more classification models!", visible=False, value="None" ) custom_weights_path = gr.Textbox(label="Custom Weights Path", visible=False, interactive=True, placeholder="./weights/my_weight.pt") custom_weights_class = gr.Textbox(label="Custom Weights Class", visible=False, interactive=True, placeholder="{1:'ocelot', 2:'cow', 3:'bear'}") load_but = gr.Button("Load Models!") load_out = gr.Text("NO MODEL LOADED!!", label="Loaded models:") def update_ui_elements(det_model): if det_model == "MegaDetectorV6": return gr.Dropdown(choices=["MDV6-yolov9-c", "MDV6-yolov9-e", "MDV6-yolov10-c", "MDV6-yolov10-e", "MDV6-rtdetr-c"], interactive=True, label="Model version", value="MDV6-yolov9e"), gr.update(visible=True) elif det_model == "MegaDetectorV5": return gr.Dropdown(choices=["a", "b"], interactive=True, label="Model version", value="a"), gr.update(visible=True) else: return gr.Dropdown(choices=["None"], interactive=True, label="Model version", value="None"), gr.update(value="None", visible=False) det_drop.change(update_ui_elements, det_drop, [det_version, clf_drop]) def toggle_textboxes(model): if model == "CustomWeights": return gr.update(visible=True), gr.update(visible=True) else: return gr.update(visible=False), gr.update(visible=False) clf_drop.change( toggle_textboxes, clf_drop, [custom_weights_path, custom_weights_class] ) with gr.Tab("Single Image Process"): with gr.Row(): with gr.Column(): sgl_in = gr.Image(type="pil") sgl_conf_sl_det = gr.Slider(0, 1, label="Detection Confidence Threshold", value=0.2) sgl_conf_sl_clf = gr.Slider(0, 1, label="Classification Confidence Threshold", value=0.7, visible=True) sgl_out = gr.Image() sgl_but = gr.Button("Detect Animals!") with gr.Tab("Folder Separation"): with gr.Row(): with gr.Column(): inp_path = gr.Textbox(label="Input path", placeholder="./data/") out_path = gr.Textbox(label="Output path", placeholder="./output/") bth_conf_fs = gr.Slider(0, 1, label="Detection Confidence Threshold", value=0.2) process_btn = gr.Button("Process Files") bth_out2 = gr.File(label="Detection Results JSON.", height=200) with gr.Column(): process_files_button = gr.Button("Separate files") process_result = gr.Text("Click on 'Separate files' once you see the JSON file", label="Separated files:") process_btn.click(batch_path_detection, inputs=[inp_path, bth_conf_fs], outputs=bth_out2) process_files_button.click(pw_utils.detection_folder_separation, inputs=[bth_out2, inp_path, out_path, bth_conf_fs], outputs=process_result) with gr.Tab("Batch Image Process"): with gr.Row(): with gr.Column(): bth_in = gr.File(label="Upload zip file.") # The timelapse checkbox is only visible when the detection model is not HerdNet chck_timelapse = gr.Checkbox(label="Generate timelapse JSON", visible=False) bth_conf_sl = gr.Slider(0, 1, label="Detection Confidence Threshold", value=0.2) bth_out = gr.File(label="Detection Results JSON.", height=200) bth_but = gr.Button("Detect Animals!") with gr.Tab("Single Video Process"): with gr.Row(): with gr.Column(): vid_in = gr.Video(label="Upload a video.") vid_conf_sl_det = gr.Slider(0, 1, label="Detection Confidence Threshold", value=0.2) vid_conf_sl_clf = gr.Slider(0, 1, label="Classification Confidence Threshold", value=0.7) vid_fr = gr.Dropdown([5, 10, 30], label="Output video framerate", value=30) vid_enc = gr.Dropdown( ["mp4v", "avc1"], label="Video encoder", info="mp4v is default, av1c is faster (needs conda install opencv)", value="mp4v" ) vid_out = gr.Video() vid_but = gr.Button("Detect Animals!") # Show timelapsed checkbox only when detection model is not HerdNet det_drop.change( lambda model: gr.update(visible=True) if "HerdNet" not in model else gr.update(visible=False), det_drop, [chck_timelapse] ) load_but.click(load_models, inputs=[det_drop, det_version, clf_drop, custom_weights_path, custom_weights_class], outputs=load_out) sgl_but.click(single_image_detection, inputs=[sgl_in, sgl_conf_sl_det, sgl_conf_sl_clf], outputs=sgl_out) bth_but.click(batch_detection, inputs=[bth_in, chck_timelapse, bth_conf_sl], outputs=bth_out) vid_but.click(video_detection, inputs=[vid_in, vid_conf_sl_det, vid_conf_sl_clf, vid_fr, vid_enc], outputs=vid_out) if __name__ == "__main__": demo.launch()