towhee-io_examples

1
{
2
 "cells": [
3
  {
4
   "cell_type": "markdown",
5
   "id": "458b4cd9",
6
   "metadata": {},
7
   "source": [
8
    "# Optimization of the Object Detection Operator with YOLO"
9
   ]
10
  },
11
  {
12
   "cell_type": "markdown",
13
   "id": "1a34f387",
14
   "metadata": {},
15
   "source": [
16
    "[Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics) is a cutting-edge, state-of-the-art (SOTA) model that builds upon the success of previous YOLO versions and introduces new features and improvements to further boost performance and flexibility. YOLOv8 is designed to be fast, accurate, and easy to use, making it an excellent choice for a wide range of object detection and tracking, instance segmentation, image classification and pose estimation tasks."
17
   ]
18
  },
19
  {
20
   "cell_type": "markdown",
21
   "id": "53dc444c",
22
   "metadata": {},
23
   "source": [
24
    "## Run with YOLO\n",
25
    "\n",
26
    "First, let's run YOLO prediction code, which will be very simple according to the [tutorial](https://docs.ultralytics.com/modes/predict/)."
27
   ]
28
  },
29
  {
30
   "cell_type": "code",
31
   "execution_count": null,
32
   "id": "4b5b2ce6",
33
   "metadata": {},
34
   "outputs": [],
35
   "source": [
36
    "from ultralytics import YOLO\n",
37
    "model = YOLO(\"yolov8n.pt\") \n",
38
    "result = model(\"https://ultralytics.com/images/bus.jpg\")[0]"
39
   ]
40
  },
41
  {
42
   "cell_type": "markdown",
43
   "id": "e01490e3",
44
   "metadata": {},
45
   "source": [
46
    "Then, we can get the information we want based on the predicted `result`, such as `boxes`, `classes` and `scores`."
47
   ]
48
  },
49
  {
50
   "cell_type": "code",
51
   "execution_count": null,
52
   "id": "5e163003",
53
   "metadata": {},
54
   "outputs": [],
55
   "source": [
56
    "boxes = [list(map(int, xyxy)) for xyxy in result.boxes.xyxy]\n",
57
    "classes = [result.names[int(i)] for i in result.boxes.cls]\n",
58
    "scores = result.boxes.conf.tolist()\n",
59
    "print(boxes, classes, scores)"
60
   ]
61
  },
62
  {
63
   "cell_type": "markdown",
64
   "id": "3c6874eb",
65
   "metadata": {},
66
   "source": [
67
    "## Develop yolo operator\n",
68
    "\n",
69
    "According to [yolov5 operator](https://towhee.io/object-detection/yolov5), we can run the pipeline for image object  detection in the following code:"
70
   ]
71
  },
72
  {
73
   "cell_type": "code",
74
   "execution_count": null,
75
   "id": "5ac79368",
76
   "metadata": {},
77
   "outputs": [],
78
   "source": [
79
    "from towhee import pipe, ops, DataCollection\n",
80
    "\n",
81
    "p = (\n",
82
    "    pipe.input('path')\n",
83
    "        .map('path', 'img', ops.image_decode.cv2_rgb())\n",
84
    "        .map('img', ('box', 'class', 'score'), ops.object_detection.yolov5())\n",
85
    "        .map(('img', 'box'), 'object', ops.image_crop(clamp=True))\n",
86
    "        .output('img', 'object', 'class')\n",
87
    ")"
88
   ]
89
  },
90
  {
91
   "cell_type": "code",
92
   "execution_count": null,
93
   "id": "c86c59d6",
94
   "metadata": {},
95
   "outputs": [],
96
   "source": [
97
    "DataCollection(p(\"https://ultralytics.com/images/bus.jpg\")).show()"
98
   ]
99
  },
100
  {
101
   "cell_type": "markdown",
102
   "id": "678b3d3a",
103
   "metadata": {},
104
   "source": [
105
    "Next, we can optimize the yolo arithmetic so as to support the yolov8 model, so we can develop yolov8 operator based on the yolov5 operator, for example, we develop the YOLOv8 class and develop the `__init__` and `__call__` function, so as to support YOLOv8 model."
106
   ]
107
  },
108
  {
109
   "cell_type": "code",
110
   "execution_count": null,
111
   "id": "e4e06fc8",
112
   "metadata": {},
113
   "outputs": [],
114
   "source": [
115
    "import numpy\n",
116
    "from towhee import register\n",
117
    "from towhee.operator import NNOperator\n",
118
    "from ultralytics import YOLO\n",
119
    "\n",
120
    "\n",
121
    "@register(name='yolov8')\n",
122
    "class YOLOv8(NNOperator):\n",
123
    "    def __init__(self, model=\"yolov8n.pt\"):\n",
124
    "        super().__init__()\n",
125
    "        self._model = YOLO(model)\n",
126
    "\n",
127
    "    def __call__(self, img: numpy.ndarray):\n",
128
    "        results = self._model(img)[0]\n",
129
    "        boxes = [list(map(int, xyxy)) for xyxy in result.boxes.xyxy]\n",
130
    "        classes = [result.names[int(i)] for i in result.boxes.cls]\n",
131
    "        scores = result.boxes.conf.tolist()\n",
132
    "        return boxes, classes, scores\n"
133
   ]
134
  },
135
  {
136
   "cell_type": "markdown",
137
   "id": "c9301040",
138
   "metadata": {},
139
   "source": [
140
    "If you compare the YOLOv8 class with the YOLOv5 class, you will see that they are very similar, the difference is the usage of the models.\n",
141
    "\n",
142
    "Then, we can also take the names `yolov8` operator to the pipeline, a similar pipeline as above, but changing the operator to YOLOv8."
143
   ]
144
  },
145
  {
146
   "cell_type": "code",
147
   "execution_count": null,
148
   "id": "a336ab8a",
149
   "metadata": {},
150
   "outputs": [],
151
   "source": [
152
    "from towhee import pipe, ops, DataCollection\n",
153
    "\n",
154
    "p = (\n",
155
    "    pipe.input('path')\n",
156
    "        .map('path', 'img', ops.image_decode.cv2_rgb())\n",
157
    "        .map('img', ('box', 'class', 'score'), ops.yolov8(model=\"yolov8n.pt\"))\n",
158
    "        .map(('img', 'box'), 'object', ops.image_crop(clamp=True))\n",
159
    "        .output('img', 'object', 'class')\n",
160
    ")"
161
   ]
162
  },
163
  {
164
   "cell_type": "code",
165
   "execution_count": null,
166
   "id": "442ebf2d",
167
   "metadata": {},
168
   "outputs": [],
169
   "source": [
170
    "DataCollection(p(\"https://ultralytics.com/images/bus.jpg\")).show()"
171
   ]
172
  },
173
  {
174
   "cell_type": "markdown",
175
   "id": "54ac8f7a",
176
   "metadata": {},
177
   "source": [
178
    "Now we have completed the development of the YOLOv8 operator, as well as adding the operator to the pipeline, and the YOLOv8 model detects one more object, the stop sign, compared to the YOLOv5 model, so it is clear that YOLOv8 is a more complete detection than YOLOv5.\n",
179
    "\n",
180
    "We can also develop the `__init__` and `__call__` methods according to how models such as YOLOv6 are used, for example, to enable different models."
181
   ]
182
  },
183
  {
184
   "cell_type": "markdown",
185
   "id": "2943fa74",
186
   "metadata": {},
187
   "source": [
188
    "## Set for training\n",
189
    "\n",
190
    "We can also train the YOLOv8 operator, according to the method of [YOLOv8 training](https://github.com/ultralytics/ultralytics/blob/dce4efce48a05e028e6ec430045431c242e52484/docs/yolov5/tutorials/train_custom_data.md), first of all, you have to manually create a \"dataset\" directory in the current directory, and then you can use COCO dataset to train the YOLOv8 model, or you can replace it with your own dataset."
191
   ]
192
  },
193
  {
194
   "cell_type": "code",
195
   "execution_count": null,
196
   "id": "9694e1e1",
197
   "metadata": {},
198
   "outputs": [],
199
   "source": [
200
    "import towhee\n",
201
    "\n",
202
    "op = towhee.ops.yolov8().get_op()"
203
   ]
204
  },
205
  {
206
   "cell_type": "code",
207
   "execution_count": null,
208
   "id": "340c963c",
209
   "metadata": {},
210
   "outputs": [],
211
   "source": [
212
    "op._model.train(data=\"coco128.yaml\", epochs=3)"
213
   ]
214
  },
215
  {
216
   "cell_type": "code",
217
   "execution_count": null,
218
   "id": "f9004cb5",
219
   "metadata": {},
220
   "outputs": [],
221
   "source": []
222
  }
223
 ],
224
 "metadata": {
225
  "kernelspec": {
226
   "display_name": "Python 3 (ipykernel)",
227
   "language": "python",
228
   "name": "python3"
229
  },
230
  "language_info": {
231
   "codemirror_mode": {
232
    "name": "ipython",
233
    "version": 3
234
   },
235
   "file_extension": ".py",
236
   "mimetype": "text/x-python",
237
   "name": "python",
238
   "nbconvert_exporter": "python",
239
   "pygments_lexer": "ipython3",
240
   "version": "3.10.12"
241
  }
242
 },
243
 "nbformat": 4,
244
 "nbformat_minor": 5
245
}
246