Skip to content

Commit

Permalink
Add Hub results.pandas() method
Browse files Browse the repository at this point in the history
New method converts results from torch tensors to pandas DataFrames with column names.

This PR may partially resolve issue #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
```
  • Loading branch information
glenn-jocher committed Apr 6, 2021
1 parent c8c8da6 commit bfba802
Showing 1 changed file with 15 additions and 3 deletions.
18 changes: 15 additions & 3 deletions models/common.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit bfba802

Please sign in to comment.