File size: 3,128 Bytes
7206ed3 e554d61 7206ed3 826067e 7206ed3 8f06817 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
from tkinter import W
import gradio as gr
from matplotlib import cm
import torch
import torchvision
import matplotlib
import PIL
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),
PIL.Image.Resampling.LANCZOS) # 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 |