@@ -147,35 +147,61 @@ def parse_hypothesis(self, results: Results) -> List[Dict]:
147
147
148
148
hypothesis_list = []
149
149
150
- box_data : Boxes
151
- for box_data in results .boxes :
152
- hypothesis = {
153
- "class_id" : int (box_data .cls ),
154
- "class_name" : self .yolo .names [int (box_data .cls )],
155
- "score" : float (box_data .conf )
156
- }
157
- hypothesis_list .append (hypothesis )
150
+ if results .boxes :
151
+ box_data : Boxes
152
+ for box_data in results .boxes :
153
+ hypothesis = {
154
+ "class_id" : int (box_data .cls ),
155
+ "class_name" : self .yolo .names [int (box_data .cls )],
156
+ "score" : float (box_data .conf )
157
+ }
158
+ hypothesis_list .append (hypothesis )
159
+
160
+ elif results .obb :
161
+ for i in range (results .obb .cls .shape [0 ]):
162
+ hypothesis = {
163
+ "class_id" : int (results .obb .cls [i ]),
164
+ "class_name" : self .yolo .names [int (results .obb .cls [i ])],
165
+ "score" : float (results .obb .conf [i ])
166
+ }
167
+ hypothesis_list .append (hypothesis )
158
168
159
169
return hypothesis_list
160
170
161
171
def parse_boxes (self , results : Results ) -> List [BoundingBox2D ]:
162
172
163
173
boxes_list = []
164
174
165
- box_data : Boxes
166
- for box_data in results .boxes :
175
+ if results .boxes :
176
+ box_data : Boxes
177
+ for box_data in results .boxes :
167
178
168
- msg = BoundingBox2D ()
179
+ msg = BoundingBox2D ()
169
180
170
- # get boxes values
171
- box = box_data .xywh [0 ]
172
- msg .center .position .x = float (box [0 ])
173
- msg .center .position .y = float (box [1 ])
174
- msg .size .x = float (box [2 ])
175
- msg .size .y = float (box [3 ])
181
+ # get boxes values
182
+ box = box_data .xywh [0 ]
183
+ msg .center .position .x = float (box [0 ])
184
+ msg .center .position .y = float (box [1 ])
185
+ msg .size .x = float (box [2 ])
186
+ msg .size .y = float (box [3 ])
176
187
177
- # append msg
178
- boxes_list .append (msg )
188
+ # append msg
189
+ boxes_list .append (msg )
190
+
191
+ elif results .obb :
192
+ for i in range (results .obb .cls .shape [0 ]):
193
+ msg = BoundingBox2D ()
194
+
195
+ # get boxes values
196
+ box = results .obb .xywhr [i ]
197
+ msg .center .position .x = float (box [0 ])
198
+ msg .center .position .y = float (box [1 ])
199
+ msg .center .theta = float (box [4 ])
200
+ msg .size .x = float (box [2 ])
201
+ msg .size .y = float (box [3 ])
202
+
203
+ # append msg
204
+ boxes_list .append (msg )
179
205
180
206
return boxes_list
181
207
@@ -246,7 +272,7 @@ def image_cb(self, msg: Image) -> None:
246
272
)
247
273
results : Results = results [0 ].cpu ()
248
274
249
- if results .boxes :
275
+ if results .boxes or results . obb :
250
276
hypothesis = self .parse_hypothesis (results )
251
277
boxes = self .parse_boxes (results )
252
278
@@ -263,7 +289,7 @@ def image_cb(self, msg: Image) -> None:
263
289
264
290
aux_msg = Detection ()
265
291
266
- if results .boxes :
292
+ if results .boxes or results . obb :
267
293
aux_msg .class_id = hypothesis [i ]["class_id" ]
268
294
aux_msg .class_name = hypothesis [i ]["class_name" ]
269
295
aux_msg .score = hypothesis [i ]["score" ]
0 commit comments