hlydecker commited on
Commit
844db40
1 Parent(s): a085433

content: add app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Marsupial Demo
2
+ import gradio as gr
3
+ import torch
4
+ import torchvision
5
+ import numpy as np
6
+ from PIL import Image
7
+
8
+ # Load Marsupial model
9
+ # TODO: Allow user selectable model?
10
+ # models = ["model_weights/md_v5a.0.0.pt","model_weights/md_v5b.0.0.pt"]
11
+ model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/marsupial_72s_lures.pt")
12
+
13
+ def yolo(im, size=640):
14
+ g = (size / max(im.size)) # gain
15
+ im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS) # resize
16
+
17
+ model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/marsupial_72s_lures.pt")
18
+
19
+ results = model(im) # inference
20
+ results.render() # updates results.imgs with boxes and labels
21
+ return Image.fromarray(results.imgs[0])
22
+
23
+ #image = gr.inputs.Image(type="pil", label="Input Image")
24
+ #chosen_model = gr.inputs.Dropdown(choices = models, value = "model_weights/md_v5a.0.0.pt",type = "value", label="Model Weight")
25
+ #size = 640
26
+
27
+ #inputs = [image, chosen_model, size]
28
+ inputs = gr.inputs.Image(type="pil", label="Input Image")
29
+ outputs = gr.outputs.Image(type="pil", label="Output Image")
30
+
31
+ title = "Marsupial"
32
+ description = "Detect and identify 72 different species of Australian wildlife using Marsupial's most detailed model"
33
+ article = "<p style='text-align: center'>This app makes predictions using a YOLOv5s model that was trained to detect and identify 72 different species of animals found in Australia in camera trap images; find out more about the project on <a href='https://github.com/Sydney-Informatics-Hub/marsupial'>GitHub</a>. This app was built by Dr Henry Lydecker and Dr Gordon MacDonald at the Sydney Informatics Hub, a Core Research Facility at the University of Sydney. Find out more about the YOLO model from the original creator, <a href='https://pjreddie.com/darknet/yolo/'>Joseph Redmon</a>. YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset and developed by Ultralytics, and includes simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"
34
+
35
+ examples = [['data/Macropod.jpg'], ['data/koala2.jpg'],['data/cat.jpg'],['data/BrushtailPossum.jpg']]
36
+ gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, theme="huggingface").launch(enable_queue=True)