Skip to content

Commit c61c4e7

Browse files
TianxiaomoTianxiaomo
authored
Merge pull request #143 from wenh06/master
add the mistakenly missing self.imgs in Yolo_dataset
2 parents 7bd1a68 + f2497b3 commit c61c4e7

File tree

3 files changed

+505
-2
lines changed

3 files changed

+505
-2
lines changed

‎dataset.py‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -260,14 +260,15 @@ def __init__(self, lable_path, cfg, train=True):
260260
truth[data[0]].append([int(j) for j in i.split(',')])
261261

262262
self.truth = truth
263+
self.imgs = list(self.truth.keys())
263264

264265
def __len__(self):
265266
return len(self.truth.keys())
266267

267268
def __getitem__(self, index):
268269
if not self.train:
269270
return self._get_val_item(index)
270-
img_path = list(self.truth.keys())[index]
271+
img_path = self.imgs[index]
271272
bboxes = np.array(self.truth.get(img_path), dtype=np.float)
272273
img_path = os.path.join(self.cfg.dataset_dir, img_path)
273274
use_mixup = self.cfg.mixup
@@ -408,7 +409,19 @@ def _get_val_item(self, index):
408409

409410

410411
def get_image_id(filename:str) -> int:
411-
"""Convert a string to a integer."""
412+
"""
413+
Convert a string to a integer.
414+
Make sure that the images and the `image_id`s are in one-one correspondence.
415+
There are already `image_id`s in annotations of the COCO dataset,
416+
in which case this function is unnecessary.
417+
For creating one's own `get_image_id` function, one can refer to
418+
https://github.com/google/automl/blob/master/efficientdet/dataset/create_pascal_tfrecord.py#L86
419+
or refer to the following code (where the filenames are like 'level1_123.jpg')
420+
>>> lv, no = os.path.splitext(os.path.basename(filename))[0].split("_")
421+
>>> lv = lv.replace("level", "")
422+
>>> no = f"{int(no):04d}"
423+
>>> return int(lv+no)
424+
"""
412425
raise NotImplementedError("Create your own 'get_image_id' function")
413426
lv, no = os.path.splitext(os.path.basename(filename))[0].split("_")
414427
lv = lv.replace("level", "")

‎tool/utils_iou.py‎

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
1+
# -*- coding: utf-8 -*-
2+
'''
3+
4+
'''
5+
import torch
6+
import os, sys
7+
from torch.nn import functional as F
8+
9+
import numpy as np
10+
from packaging import version
11+
12+
13+
__all__ = [
14+
"bboxes_iou",
15+
"bboxes_giou",
16+
"bboxes_diou",
17+
"bboxes_ciou",
18+
]
19+
20+
21+
if version.parse(torch.__version__) >= version.parse('1.5.0'):
22+
def _true_divide(dividend, divisor):
23+
return torch.true_divide(dividend, divisor)
24+
else:
25+
def _true_divide(dividend, divisor):
26+
return dividend / divisor
27+
28+
def bboxes_iou(bboxes_a, bboxes_b, fmt='voc', iou_type='iou'):
29+
"""Calculate the Intersection of Unions (IoUs) between bounding boxes.
30+
IoU is calculated as a ratio of area of the intersection
31+
and area of the union.
32+
33+
Args:
34+
bbox_a (array): An array whose shape is :math:`(N, 4)`.
35+
:math:`N` is the number of bounding boxes.
36+
The dtype should be :obj:`numpy.float32`.
37+
bbox_b (array): An array similar to :obj:`bbox_a`,
38+
whose shape is :math:`(K, 4)`.
39+
The dtype should be :obj:`numpy.float32`.
40+
Returns:
41+
array:
42+
An array whose shape is :math:`(N, K)`. \
43+
An element at index :math:`(n, k)` contains IoUs between \
44+
:math:`n` th bounding box in :obj:`bbox_a` and :math:`k` th bounding \
45+
box in :obj:`bbox_b`.
46+
47+
from: https://github.com/chainer/chainercv
48+
"""
49+
if bboxes_a.shape[1] != 4 or bboxes_b.shape[1] != 4:
50+
raise IndexError
51+
52+
N, K = bboxes_a.shape[0], bboxes_b.shape[0]
53+
54+
if fmt.lower() == 'voc': # xmin, ymin, xmax, ymax
55+
# top left
56+
tl_intersect = torch.max(
57+
bboxes_a[:, np.newaxis, :2],
58+
bboxes_b[:, :2]
59+
) # of shape `(N,K,2)`
60+
# bottom right
61+
br_intersect = torch.min(
62+
bboxes_a[:, np.newaxis, 2:],
63+
bboxes_b[:, 2:]
64+
)
65+
bb_a = bboxes_a[:, 2:] - bboxes_a[:, :2]
66+
bb_b = bboxes_b[:, 2:] - bboxes_b[:, :2]
67+
# bb_* can also be seen vectors representing box_width, box_height
68+
elif fmt.lower() == 'yolo': # xcen, ycen, w, h
69+
# top left
70+
tl_intersect = torch.max(
71+
bboxes_a[:, np.newaxis, :2] - bboxes_a[:, np.newaxis, 2:] / 2,
72+
bboxes_b[:, :2] - bboxes_b[:, 2:] / 2
73+
)
74+
# bottom right
75+
br_intersect = torch.min(
76+
bboxes_a[:, np.newaxis, :2] + bboxes_a[:, np.newaxis, 2:] / 2,
77+
bboxes_b[:, :2] + bboxes_b[:, 2:] / 2
78+
)
79+
bb_a = bboxes_a[:, 2:]
80+
bb_b = bboxes_b[:, 2:]
81+
elif fmt.lower() == 'coco': # xmin, ymin, w, h
82+
# top left
83+
tl_intersect = torch.max(
84+
bboxes_a[:, np.newaxis, :2],
85+
bboxes_b[:, :2]
86+
)
87+
# bottom right
88+
br_intersect = torch.min(
89+
bboxes_a[:, np.newaxis, :2] + bboxes_a[:, np.newaxis, 2:],
90+
bboxes_b[:, :2] + bboxes_b[:, 2:]
91+
)
92+
bb_a = bboxes_a[:, 2:]
93+
bb_b = bboxes_b[:, 2:]
94+
95+
area_a = torch.prod(bb_a, 1)
96+
area_b = torch.prod(bb_b, 1)
97+
98+
# torch.prod(input, dim, keepdim=False, dtype=None) → Tensor
99+
# Returns the product of each row of the input tensor in the given dimension dim
100+
# if tl, br does not form a nondegenerate squre, then the corr. element in the `prod` would be 0
101+
en = (tl_intersect < br_intersect).type(tl_intersect.type()).prod(dim=2) # shape `(N,K,2)` ---> shape `(N,K)`
102+
103+
area_intersect = torch.prod(br_intersect - tl_intersect, 2) * en # * ((tl < br).all())
104+
area_union = (area_a[:, np.newaxis] + area_b - area_intersect)
105+
106+
iou = _true_divide(area_intersect, area_union)
107+
108+
if iou_type.lower() == 'iou':
109+
return iou
110+
111+
if fmt.lower() == 'voc': # xmin, ymin, xmax, ymax
112+
# top left
113+
tl_union = torch.min(
114+
bboxes_a[:, np.newaxis, :2],
115+
bboxes_b[:, :2]
116+
) # of shape `(N,K,2)`
117+
# bottom right
118+
br_union = torch.max(
119+
bboxes_a[:, np.newaxis, 2:],
120+
bboxes_b[:, 2:]
121+
)
122+
elif fmt.lower() == 'yolo': # xcen, ycen, w, h
123+
# top left
124+
tl_union = torch.min(
125+
bboxes_a[:, np.newaxis, :2] - bboxes_a[:, np.newaxis, 2:] / 2,
126+
bboxes_b[:, :2] - bboxes_b[:, 2:] / 2
127+
)
128+
# bottom right
129+
br_union = torch.max(
130+
bboxes_a[:, np.newaxis, :2] + bboxes_a[:, np.newaxis, 2:] / 2,
131+
bboxes_b[:, :2] + bboxes_b[:, 2:] / 2
132+
)
133+
elif fmt.lower() == 'coco': # xmin, ymin, w, h
134+
# top left
135+
tl_union = torch.min(
136+
bboxes_a[:, np.newaxis, :2],
137+
bboxes_b[:, :2]
138+
)
139+
# bottom right
140+
br_union = torch.max(
141+
bboxes_a[:, np.newaxis, :2] + bboxes_a[:, np.newaxis, 2:],
142+
bboxes_b[:, :2] + bboxes_b[:, 2:]
143+
)
144+
145+
# c for covering, of shape `(N,K,2)`
146+
# the last dim is box width, box hight
147+
bboxes_c = br_union - tl_union
148+
149+
area_covering = torch.prod(bboxes_c, 2) # shape `(N,K)`
150+
151+
giou = iou - _true_divide(area_covering - area_union, area_covering)
152+
153+
if iou_type.lower() == 'giou':
154+
return giou
155+
156+
if fmt.lower() == 'voc': # xmin, ymin, xmax, ymax
157+
centre_a = (bboxes_a[..., 2 :] + bboxes_a[..., : 2]) / 2
158+
centre_b = (bboxes_b[..., 2 :] + bboxes_b[..., : 2]) / 2
159+
elif fmt.lower() == 'yolo': # xcen, ycen, w, h
160+
centre_a = bboxes_a[..., : 2]
161+
centre_b = bboxes_b[..., : 2]
162+
elif fmt.lower() == 'coco': # xmin, ymin, w, h
163+
centre_a = bboxes_a[..., 2 :] + bboxes_a[..., : 2]/2
164+
centre_b = bboxes_b[..., 2 :] + bboxes_b[..., : 2]/2
165+
166+
centre_dist = torch.norm(centre_a[:, np.newaxis] - centre_b, p='fro', dim=2)
167+
diag_len = torch.norm(bboxes_c, p='fro', dim=2)
168+
169+
diou = iou - _true_divide(centre_dist.pow(2), diag_len.pow(2))
170+
171+
if iou_type.lower() == 'diou':
172+
return diou
173+
174+
# bb_a of shape `(N,2)`, bb_b of shape `(K,2)`
175+
v = torch.einsum('nm,km->nk', bb_a, bb_b)
176+
v = _true_divide(v, (torch.norm(bb_a, p='fro', dim=1)[:,np.newaxis] * torch.norm(bb_b, p='fro', dim=1)))
177+
# avoid nan for torch.acos near \pm 1
178+
# https://github.com/pytorch/pytorch/issues/8069
179+
eps = 1e-7
180+
v = torch.clamp(v, -1+eps, 1-eps)
181+
v = (_true_divide(2*torch.acos(v), np.pi)).pow(2)
182+
with torch.no_grad():
183+
alpha = (_true_divide(v, 1-iou+v)) * ((iou>=0.5).type(iou.type()))
184+
185+
ciou = diou - alpha * v
186+
187+
if iou_type.lower() == 'ciou':
188+
return ciou
189+
190+
191+
def bboxes_giou(bboxes_a, bboxes_b, fmt='voc'):
192+
return bboxes_iou(bboxes_a, bboxes_b, fmt, 'giou')
193+
194+
195+
def bboxes_diou(bboxes_a, bboxes_b, fmt='voc'):
196+
return bboxes_iou(bboxes_a, bboxes_b, fmt, 'diou')
197+
198+
199+
def bboxes_ciou(bboxes_a, bboxes_b, fmt='voc'):
200+
return bboxes_iou(bboxes_a, bboxes_b, fmt, 'ciou')

0 commit comments

Comments
 (0)