From bfba802964322336d6520ed90a7ab9f9eac162c0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 7 Apr 2021 01:50:38 +0200 Subject: [PATCH] Add Hub results.pandas() method New method converts results from torch tensors to pandas DataFrames with column names. This PR may partially resolve issue https://github.com/ultralytics/yolov5/issues/2703 ```python print(results.pandas().xyxy[0]) xmin ymin xmax ymax confidence class name 0 57.068970 391.770599 241.383545 905.797852 0.868964 0.0 person 1 667.661255 399.303589 810.000000 881.396667 0.851888 0.0 person 2 222.878387 414.774231 343.804474 857.825073 0.838376 0.0 person 3 4.205386 234.447678 803.739136 750.023376 0.658006 5.0 bus 4 0.000000 550.596008 76.681190 878.669922 0.450596 0.0 person ``` --- models/common.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/models/common.py b/models/common.py index 4fd1a8159c64..c7571c202231 100644 --- a/models/common.py +++ b/models/common.py @@ -1,9 +1,11 @@ # YOLOv5 common modules import math +from copy import copy from pathlib import Path import numpy as np +import pandas as pd import requests import torch import torch.nn as nn @@ -347,17 +349,27 @@ def render(self): self.display(render=True) # render results return self.imgs - def __len__(self): - return self.n + def pandas(self): + # return detections as pandas DataFrames + new = copy(self) # do not replace self + ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns + cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns + for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]): + a = [[x + [self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # updated attribute + setattr(new, k, [pd.DataFrame(x, columns=c) for x in a]) + return new def tolist(self): # return a list of Detections objects, i.e. 'for result in results.tolist():' - x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)] + x = [Detections([self.imgs[i]], [self.pred[i]], self.names, self.s) for i in range(self.n)] for d in x: for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']: setattr(d, k, getattr(d, k)[0]) # pop out of list return x + def __len__(self): + return self.n + class Classify(nn.Module): # Classification head, i.e. x(b,c1,20,20) to x(b,c2)