diff --git a/README.md b/README.md
index 4fb3e03..a756e75 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,87 @@
-# yolo
+# Object Detection with Yolov5
+
+*author: shiyu22*
+
+
+
+
+
+
+
+### 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.
+
+
+
+
+
+
+
+### Code Example
+
+
+
+Writing the pipeline in the simplified way
+
+```Python
+import towhee
+
+towhee.glob('./test.jpg') \
+ .image_decode() \
+ .object_detection.yolov5() \
+ .show()
+```
+
+
+
+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()
+```
+
+
+
+
+
+
+
+
+
+## Factory Constructor
+
+Create the operator via the following factory method
+
+***object_detection.yolov5()***
+
+
+
+
+
+
+
+### 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.
diff --git a/__init__.py b/__init__.py
index 64db8e3..2700aa5 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,5 +1,4 @@
-from yolov5.yolov5 import Yolov5
-
+from .yolov5 import Yolov5
def yolov5():
return Yolov5()
\ No newline at end of file
diff --git a/results1.png b/results1.png
new file mode 100644
index 0000000..d21fd20
Binary files /dev/null and b/results1.png differ
diff --git a/results2.png b/results2.png
new file mode 100644
index 0000000..7a8b970
Binary files /dev/null and b/results2.png differ
diff --git a/test.png b/test.png
new file mode 100755
index 0000000..c3ca0be
Binary files /dev/null and b/test.png differ
diff --git a/yolov5.py b/yolov5.py
index b60e2d8..9e927f9 100644
--- a/yolov5.py
+++ b/yolov5.py
@@ -1,18 +1,21 @@
import torch
import numpy
+
+from towhee import register
from towhee.operator import NNOperator
-@register(outputschema=['boxes', 'classes', 'scores'])
+@register(output_schema=['boxes', 'classes', 'scores'])
class Yolov5(NNOperator):
- def __init__(self, model_name):
+ def __init__(self, model_name: str ='yolov5s'):
super().__init__()
self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True)
def __call__(self, img: numpy.ndarray):
# Get object detection results with YOLOv5 model
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)
scores = list(results.pandas().xyxy[0].confidence)
return boxes, classes, scores