logo
Browse Source

Update yolov5 README

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 3 years ago
parent
commit
bb2d3fce7e
  1. 87
      README.md
  2. 3
      __init__.py
  3. BIN
      results1.png
  4. BIN
      results2.png
  5. BIN
      test.png
  6. 9
      yolov5.py

87
README.md

@ -1,2 +1,87 @@
# yolo
# Object Detection with Yolov5
*author: shiyu22*
<br />
### Description
**Object Detection** is a computer vision technique that locates and identifies people, items, or other objects in an image. Object detection has applications in many areas of computer vision, including image retrieval, image annotation, vehicle counting, object tracking, etc.
This operator uses [PyTorch.yolov5](https://pytorch.org/hub/ultralytics_yolov5/) to detect the object.
<br />
### Code Example
Writing the pipeline in the simplified way
```Python
import towhee
towhee.glob('./test.jpg') \
.image_decode() \
.object_detection.yolov5() \
.show()
```
<img src="./results1.png" alt="results1" height="40px"/>
Writing the same pipeline with explicitly specified inputs and outputs
```Python
import towhee
towhee.glob['path']('./test.png') \
.image_decode['path','img']() \
.object_detection.yolov5['img', ('box', 'class', 'score')]() \
.select('box', 'class', 'score') \
.show()
```
<img src="./results2.png" alt="results1" height="90px"/>
<br />
## Factory Constructor
Create the operator via the following factory method
***object_detection.yolov5()***
<br />
### Interface
The operator takes an image as input. It first detects the objects appeared in the image, and gives the bounding box of each object.
**Parameters:**
**img**: numpy.ndarray
​ Image data in ndarray format.
**Return**: List[(int, int, int, int)], List[str], List[float]
The return value is a tuple of (boxes, classes, scores). The *boxes* is a list of bounding boxes. Each bounding box is represented by the top-left and the bottom right points, i.e. (x1, y1, x2, y2). The *classes* is a list of prediction labels. The *scores* is a list of the confidence scores.

3
__init__.py

@ -1,5 +1,4 @@
from yolov5.yolov5 import Yolov5
from .yolov5 import Yolov5
def yolov5(): def yolov5():
return Yolov5() return Yolov5()

BIN
results1.png

Binary file not shown.

After

Width:  |  Height:  |  Size: 18 KiB

BIN
results2.png

Binary file not shown.

After

Width:  |  Height:  |  Size: 54 KiB

BIN
test.png

Binary file not shown.

After

Width:  |  Height:  |  Size: 300 KiB

9
yolov5.py

@ -1,10 +1,12 @@
import torch import torch
import numpy import numpy
from towhee import register
from towhee.operator import NNOperator from towhee.operator import NNOperator
@register(outputschema=['boxes', 'classes', 'scores'])
@register(output_schema=['boxes', 'classes', 'scores'])
class Yolov5(NNOperator): class Yolov5(NNOperator):
def __init__(self, model_name):
def __init__(self, model_name: str ='yolov5s'):
super().__init__() super().__init__()
self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True) self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True)
@ -12,7 +14,8 @@ class Yolov5(NNOperator):
# Get object detection results with YOLOv5 model # Get object detection results with YOLOv5 model
results = self._model(img) results = self._model(img)
boxes = results.xyxy[0]
boxes = [re[0:3] for re in results.xyxy[0]]
boxes = [list(map(int, box)) for box in boxes]
classes = list(results.pandas().xyxy[0].name) classes = list(results.pandas().xyxy[0].name)
scores = list(results.pandas().xyxy[0].confidence) scores = list(results.pandas().xyxy[0].confidence)
return boxes, classes, scores return boxes, classes, scores

Loading…
Cancel
Save