File size: 2,554 Bytes
8960ee7 06c6438 8960ee7 658151c 06c6438 140414b 8960ee7 140414b 8960ee7 06c6438 658151c 0212b65 140414b b86d4ed 140414b 658151c 140414b 8960ee7 140414b 8960ee7 658151c |
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 |
# MegaDetector v5 and DLC Demo
import gradio as gr
import torch
import torchvision
import numpy as np
from PIL import Image
#script load
import json
import os
import numpy as np
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
from dlclive import DLCLive, Processor
from numpy import savetxt
# Load MegaDetector v5a model
model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/md_v5a.0.0.pt")
def yolo(im, size=640):
g = (size / max(im.size)) # gain
im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS) # resize
model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/md_v5a.0.0.pt")
results = model(im) # inference
results.render() # updates results.imgs with boxes and labels
return Image.fromarray(results.imgs[0])
def dlclive_pose(model, crop_np, crop, fname, index,dlc_proc):
model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/md_v5a.0.0.pt")
dlc_live = DLCLive(model, processor=dlc_proc)
dlc_live.init_inference(crop_np)
keypts = dlc_live.get_pose(crop_np)
savetxt(str(fname)+ '_' + str(index) + '.csv' , keypts, delimiter=',')
xpose = []
ypose = []
for key in keypts[:,2]:
# if key > 0.05: # which value do we need here?
i = np.where(keypts[:,2]==key)
xpose.append(keypts[i,0])
ypose.append(keypts[i,1])
plt.imshow(crop)
plt.scatter(xpose[:], ypose[:], 40, color='cyan')
plt.savefig(str(fname)+ '_' + str(index) + '.png')
plt.show()
plt.clf()
dlc_proc = Processor()
#Layouts and descriptions
title = "MegaDetector and DeepLabcutLive"
description = "Interact with MegaDetector and DeeplabCutLive for pose estimation"
article = "<p style='text-align: center'>This app uses MegaDetector YOLOv5x6 model that was trained to detect animals, humans, and vehicles in camera trap images; find out more about the project on <a href='https://github.com/microsoft/CameraTraps'>GitHub</a>. We have also integrated DeepLabCut Live for pose estimation <a href='https://github.com/DeepLabCut/DeepLabCut-live'></a>.</p>"
# input image and output image
inputs = gr.inputs.Image(type="pil", label="Input Image")
outputs = gr.outputs.Image(type="pil", label="Output Image")
#data images stored
examples = [['data/owl.jpg'], ['data/snake.jpg'],['data/beluga.jpg'],['data/rhino.jpg']]
gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, theme="huggingface").launch(enable_queue=True) |