Spaces:
Runtime error
Runtime error
from tkinter import W | |
import gradio as gr | |
from matplotlib import cm | |
import torch | |
import torchvision | |
import matplotlib | |
from PIL import Image, ImageColor, ImageFont, ImageDraw | |
import numpy as np | |
import math | |
import yaml | |
import pdb | |
############################################ | |
# Predict detections with MegaDetector v5a model | |
def predict_md(im, | |
megadetector_model, #Megadet_Models[mega_model_input] | |
size=640): | |
# resize image | |
g = (size / max(im.size)) # multipl factor to make max size of the image equal to input size | |
im = im.resize((int(x * g) for x in im.size), | |
Image.ANTIALIAS) # resize | |
# device | |
if torch.cuda.is_available(): | |
md_device = torch.device('cuda') | |
else: | |
md_device = torch.device('cpu') | |
# megadetector | |
MD_model = torch.hub.load('ultralytics/yolov5', # repo_or_dir | |
'custom', #model | |
megadetector_model, # args for callable model | |
force_reload=True, | |
device=md_device) | |
# send model to gpu if possible | |
if (md_device == torch.device('cuda')): | |
print('Sending model to GPU') | |
MD_model.to(md_device) | |
## detect objects | |
results = MD_model(im) # inference # vars(results).keys()= dict_keys(['imgs', 'pred', 'names', 'files', 'times', 'xyxy', 'xywh', 'xyxyn', 'xywhn', 'n', 't', 's']) | |
return results | |
########################################## | |
def crop_animal_detections(img_in, | |
yolo_results, | |
likelihood_th): | |
## Extract animal crops | |
list_labels_as_str = [i for i in yolo_results.names.values()] # ['animal', 'person', 'vehicle'] | |
list_np_animal_crops = [] | |
# image to crop (scale as input for megadetector) | |
img_in = img_in.resize((yolo_results.ims[0].shape[1], | |
yolo_results.ims[0].shape[0])) | |
# for every detection in the img | |
for det_array in yolo_results.xyxy: | |
# for every detection | |
for j in range(det_array.shape[0]): | |
# compute coords around bbox rounded to the nearest integer (for pasting later) | |
xmin_rd = int(math.floor(det_array[j,0])) # int() should suffice? | |
ymin_rd = int(math.floor(det_array[j,1])) | |
xmax_rd = int(math.ceil(det_array[j,2])) | |
ymax_rd = int(math.ceil(det_array[j,3])) | |
pred_llk = det_array[j,4] | |
pred_label = det_array[j,5] | |
# keep animal crops above threshold | |
if (pred_label == list_labels_as_str.index('animal')) and \ | |
(pred_llk >= likelihood_th): | |
area = (xmin_rd, ymin_rd, xmax_rd, ymax_rd) | |
#pdb.set_trace() | |
crop = img_in.crop(area) #Image.fromarray(img_in).crop(area) | |
crop_np = np.asarray(crop) | |
# add to list | |
list_np_animal_crops.append(crop_np) | |
return list_np_animal_crops |