import random import numpy as np import gradio as gr from huggingface_hub import from_pretrained_fastai from PIL import Image from groundingdino.util.inference import load_model from groundingdino.util.inference import predict as grounding_dino_predict import groundingdino.datasets.transforms as T import torch from torchvision.ops import box_convert from torchvision.transforms.functional import to_tensor from torchvision.transforms import GaussianBlur import time from Ambrosia import pre_process_image # Automatically set device based on availability DEVICE = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {DEVICE}") PROMPT = "bug" # Define a custom transform for Gaussian blur def gaussian_blur(x, p=0.5, kernel_size_min=3, kernel_size_max=20, sigma_min=0.1, sigma_max=3): if x.ndim == 4: for i in range(x.shape[0]): if random.random() < p: kernel_size = random.randrange(kernel_size_min, kernel_size_max + 1, 2) sigma = random.uniform(sigma_min, sigma_max) x[i] = GaussianBlur(kernel_size=kernel_size, sigma=sigma)(x[i]) return x # Custom Label Function def custom_label_func(fpath): # this directs the labels to be 2 levels up from the image folder label = fpath.parents[2].name return label # this function only describes how much a singular value in al ist stands out. # if all values in the lsit are high or low this is 1 # the smaller the proportiopn of number of disimilar vlaues are to other more similar values the lower this number # the larger the gap between the dissimilar numbers and the simialr number the smaller this number # only able to interpret probabilities or values between 0 and 1 # this function outputs an estimate an inverse of the classification confidence based on the probabilities of all the classes. # the wedge threshold splits the data on a threshold with a magnitude of a positive int to force a ledge/peak in the data def unkown_prob_calc(probs, wedge_threshold, wedge_magnitude=1, wedge='strict'): if wedge =='strict': increase_var = (1/(wedge_magnitude)) decrease_var = (wedge_magnitude) if wedge =='dynamic': # this allows pointsthat are furhter from the threshold ot be moved less and points clsoer to be moved more increase_var = (1/(wedge_magnitude*((1-np.abs(probs-wedge_threshold))))) decrease_var = (wedge_magnitude*((1-np.abs(probs-wedge_threshold)))) else: print("Error: use 'strict' (default) or 'dynamic' as options for the wedge parameter!") probs = np.where(probs>=wedge_threshold , probs**increase_var, probs) probs = np.where(probs<=wedge_threshold , probs**decrease_var, probs) diff_matrix = np.abs(probs[:, np.newaxis] - probs) diff_matrix_sum = np.sum(diff_matrix) probs_sum = np.sum(probs) class_val = (diff_matrix_sum/probs_sum) max_class_val = ((len(probs)-1)*2) kown_prob = class_val/max_class_val unknown_prob = 1-kown_prob return(unknown_prob) def load_image(image_source): transform = T.Compose( [ T.RandomResize([800], max_size=1333), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) image_source = image_source.convert("RGB") image_transformed, _ = transform(image_source, None) return image_transformed # load object detection model od_model = load_model( model_checkpoint_path="groundingdino_swint_ogc.pth", model_config_path="GroundingDINO_SwinT_OGC.cfg.py", device=DEVICE) print("Object detection model loaded") def detect_objects(og_image, model=od_model, prompt="bug . insect", device="cpu"): TEXT_PROMPT = prompt BOX_TRESHOLD = 0.35 TEXT_TRESHOLD = 0.25 DEVICE = device # cuda or cpu # Convert numpy array to PIL Image if needed if isinstance(og_image, np.ndarray): og_image_obj = Image.fromarray(og_image) else: og_image_obj = og_image # Assuming og_image is already a PIL Image # Transform the image image_transformed = load_image(image_source = og_image_obj) # Your model prediction code here... boxes, logits, phrases = grounding_dino_predict( model=model, image=image_transformed, caption=TEXT_PROMPT, box_threshold=BOX_TRESHOLD, text_threshold=TEXT_TRESHOLD, device=DEVICE) # Use og_image_obj directly for further processing height, width = og_image_obj.size boxes_norm = boxes * torch.Tensor([height, width, height, width]) xyxy = box_convert( boxes=boxes_norm, in_fmt="cxcywh", out_fmt="xyxy").numpy() img_lst = [] for i in range(len(boxes_norm)): crop_img = og_image_obj.crop((xyxy[i])) img_lst.append(crop_img) return (img_lst) # load beetle classifier model repo_id="ChristopherMarais/beetle-model-mini" bc_model = from_pretrained_fastai(repo_id) bc_model.to(DEVICE) # get class names labels = np.append(np.array(bc_model.dls.vocab), "Unknown") # Replace some names used in the classifier # Check if the element was found to prevent errors # The target value you're looking for target = "Scolotodes_schwarzi" # Finding the index using np.where indices = np.where(labels == target) # Extracting the first occurrence, if found if indices[0].size > 0: idx = indices[0][0] print(f"Index of {target}: {idx}") else: print(f"{target} not found in the array.") # Replace occurence if idx != -1: labels[idx] = "Scolytodes_glaber" print("Classification model loaded") def predict_beetle(img): print("Detecting & classifying beetles...") start_time = time.perf_counter() # record how long it processes # Split image into smaller images of detected objects image_lst = detect_objects(og_image=img, model=od_model, prompt=PROMPT, device=DEVICE) # pre_process = pre_process_image(manual_thresh_buffer=0.15, image = img) # use image_dir if directory of image used # pre_process.segment(cluster_num=2, # image_edge_buffer=50) # image_lst = pre_process.col_image_lst print("Objects detected") end_time = time.perf_counter() processing_time = end_time - start_time print(f"Processing duration: {processing_time} seconds") # get predictions for all segments conf_dict_lst = [] output_lst = [] img_cnt = len(image_lst) for i in range(0,img_cnt): prob_ar = np.array(bc_model.predict(image_lst[i])[2].to(DEVICE).cpu()) unkown_prob = unkown_prob_calc(probs=prob_ar, wedge_threshold=0.85, wedge_magnitude=5, wedge='dynamic') prob_ar = np.append(prob_ar, unkown_prob) prob_ar = np.around(prob_ar*100, decimals=1) # only show the top 5 predictions # Sorting the dictionary by value in descending order and taking the top items top_num = 3 conf_dict = {labels[i]: float(prob_ar[i]) for i in range(len(prob_ar))} print(conf_dict) conf_dict = dict(sorted(conf_dict.items(), key=lambda item: item[1], reverse=True)[:top_num]) conf_dict_lst.append(str(conf_dict)[1:-1]) # remove dictionary brackets result = list(zip(image_lst, conf_dict_lst)) print(f"Beetle classified - {i}") # record how long classification takes end_time = time.perf_counter() processing_time = end_time - start_time print(f"Processing duration: {processing_time} seconds") return(result) # gradio app css = """ button { width: auto; /* Set your desired width */ } """ with gr.Blocks(css=css) as demo: gr.Markdown("

Bark Beetle Classification

") gr.Markdown("

Note this instance of the classifier is for demonstration only and runs on CPU, not on GPU. If you are interested in testing the model, contact us, and we will switch it to its full capacity in an instant.

") with gr.Column(variant="panel"): with gr.Row(variant="compact"): inputs = gr.Image() # Use the `full_width` parameter directly btn = gr.Button("Classify") # Set the gallery layout and height directly in the constructor gallery = gr.Gallery(label="Show images", show_label=True, elem_id="gallery", columns=8, height="auto") btn.click(predict_beetle, inputs, gallery) demo.launch(debug=True, show_error=True)