Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
41464e18
MegEngine
项目概览
MegEngine 天元
/
MegEngine
大约 1 年 前同步成功
通知
396
Star
4705
Fork
582
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
MegEngine
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
41464e18
编写于
4月 26, 2020
作者:
M
Megvii Engine Team
提交者:
Xinran Xu
5月 06, 2020
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
feat(mge/data): voc dataset supports detection
GitOrigin-RevId: f78bef3cd2ba895c2ec0ca1e0132737cf5a5cbd5
上级
26738d99
变更
1
隐藏空白更改
内联
并排
Showing
1 changed file
with
43 addition
and
31 deletion
+43
-31
python_module/megengine/data/dataset/vision/voc.py
python_module/megengine/data/dataset/vision/voc.py
+43
-31
未找到文件。
python_module/megengine/data/dataset/vision/voc.py
浏览文件 @
41464e18
...
@@ -29,43 +29,47 @@ class PascalVOC(VisionDataset):
...
@@ -29,43 +29,47 @@ class PascalVOC(VisionDataset):
supported_order
=
(
supported_order
=
(
"image"
,
"image"
,
#
"boxes",
"boxes"
,
#
"boxes_category",
"boxes_category"
,
"mask"
,
"mask"
,
"info"
,
"info"
,
)
)
def
__init__
(
self
,
root
,
image_set
,
*
,
order
=
None
):
def
__init__
(
self
,
root
,
image_set
,
*
,
order
=
None
):
if
(
"boxes"
in
order
or
"boxes_category"
in
order
)
and
"mask"
in
order
:
raise
ValueError
(
"PascalVOC only supports boxes & boxes_category or mask, not both."
)
super
().
__init__
(
root
,
order
=
order
,
supported_order
=
self
.
supported_order
)
super
().
__init__
(
root
,
order
=
order
,
supported_order
=
self
.
supported_order
)
voc_root
=
self
.
root
if
not
os
.
path
.
isdir
(
self
.
root
):
if
not
os
.
path
.
isdir
(
voc_root
):
raise
RuntimeError
(
"Dataset not found or corrupted."
)
raise
RuntimeError
(
"Dataset not found or corrupted."
)
self
.
image_set
=
image_set
self
.
image_set
=
image_set
image_dir
=
os
.
path
.
join
(
voc_root
,
"JPEGImages"
)
image_dir
=
os
.
path
.
join
(
self
.
root
,
"JPEGImages"
)
# for segmentation
if
"boxes"
in
order
or
"boxes_category"
in
order
:
if
"aug"
in
image_set
:
annotation_dir
=
os
.
path
.
join
(
self
.
root
,
'Annotations'
)
mask_dir
=
os
.
path
.
join
(
voc_root
,
"SegmentationClass_aug"
)
splitdet_dir
=
os
.
path
.
join
(
self
.
root
,
"ImageSets/Main"
)
split_f
=
os
.
path
.
join
(
splitdet_dir
,
image_set
.
rstrip
(
"
\n
"
)
+
".txt"
)
with
open
(
os
.
path
.
join
(
split_f
),
"r"
)
as
f
:
self
.
file_names
=
[
x
.
strip
()
for
x
in
f
.
readlines
()]
self
.
images
=
[
os
.
path
.
join
(
image_dir
,
x
+
".jpg"
)
for
x
in
self
.
file_names
]
self
.
annotations
=
[
os
.
path
.
join
(
annotation_dir
,
x
+
".xml"
)
for
x
in
self
.
file_names
]
assert
len
(
self
.
images
)
==
len
(
self
.
annotations
)
elif
"mask"
in
order
:
if
"aug"
in
image_set
:
mask_dir
=
os
.
path
.
join
(
self
.
root
,
"SegmentationClass_aug"
)
else
:
mask_dir
=
os
.
path
.
join
(
self
.
root
,
"SegmentationClass"
)
splitmask_dir
=
os
.
path
.
join
(
self
.
root
,
"ImageSets/Segmentation"
)
split_f
=
os
.
path
.
join
(
splitmask_dir
,
image_set
.
rstrip
(
"
\n
"
)
+
".txt"
)
with
open
(
os
.
path
.
join
(
split_f
),
"r"
)
as
f
:
self
.
file_names
=
[
x
.
strip
()
for
x
in
f
.
readlines
()]
self
.
images
=
[
os
.
path
.
join
(
image_dir
,
x
+
".jpg"
)
for
x
in
self
.
file_names
]
self
.
masks
=
[
os
.
path
.
join
(
mask_dir
,
x
+
".png"
)
for
x
in
self
.
file_names
]
assert
len
(
self
.
images
)
==
len
(
self
.
masks
)
else
:
else
:
mask_dir
=
os
.
path
.
join
(
voc_root
,
"SegmentationClass"
)
raise
NotImplementedError
splitmask_dir
=
os
.
path
.
join
(
voc_root
,
"ImageSets/Segmentation"
)
split_f
=
os
.
path
.
join
(
splitmask_dir
,
image_set
.
rstrip
(
"
\n
"
)
+
".txt"
)
with
open
(
os
.
path
.
join
(
split_f
),
"r"
)
as
f
:
self
.
file_names
=
[
x
.
strip
()
for
x
in
f
.
readlines
()]
self
.
images
=
[
os
.
path
.
join
(
image_dir
,
x
+
".jpg"
)
for
x
in
self
.
file_names
]
self
.
masks
=
[
os
.
path
.
join
(
mask_dir
,
x
+
".png"
)
for
x
in
self
.
file_names
]
# TODO: for detection
# splitdet_dir = os.path.join(voc_root, "ImageSets/Main")
# split_f = os.path.join(splitdet_dir, image_set.rstrip("\n") + ".txt")
# with open(os.path.join(split_f), "r") as f:
# self.file_names = [x.strip() for x in f.readlines()]
# self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names]
# self.annotations = [os.path.join(annotation_dir, x + ".xml") for x in self.file_names]
# assert (len(self.images) == len(self.masks)) and (len(self.images) == len(self.annotations))
def
__getitem__
(
self
,
index
):
def
__getitem__
(
self
,
index
):
target
=
[]
target
=
[]
...
@@ -73,6 +77,19 @@ class PascalVOC(VisionDataset):
...
@@ -73,6 +77,19 @@ class PascalVOC(VisionDataset):
if
k
==
"image"
:
if
k
==
"image"
:
image
=
cv2
.
imread
(
self
.
images
[
index
],
cv2
.
IMREAD_COLOR
)
image
=
cv2
.
imread
(
self
.
images
[
index
],
cv2
.
IMREAD_COLOR
)
target
.
append
(
image
)
target
.
append
(
image
)
elif
k
==
"boxes"
:
anno
=
self
.
parse_voc_xml
(
ET
.
parse
(
self
.
annotations
[
index
]).
getroot
())
boxes
=
[
obj
[
"bndbox"
]
for
obj
in
anno
[
"annotation"
][
"object"
]]
# boxes type xyxy
boxes
=
[(
bb
[
'xmin'
],
bb
[
'ymin'
],
bb
[
'xmax'
],
bb
[
'ymax'
])
for
bb
in
boxes
]
boxes
=
np
.
array
(
boxes
,
dtype
=
np
.
float32
).
reshape
(
-
1
,
4
)
target
.
append
(
boxes
)
elif
k
==
"boxes_category"
:
anno
=
self
.
parse_voc_xml
(
ET
.
parse
(
self
.
annotations
[
index
]).
getroot
())
boxes_category
=
[
obj
[
"name"
]
for
obj
in
anno
[
"annotation"
][
"object"
]]
boxes_category
=
[
self
.
class_names
.
index
(
bc
)
-
1
for
bc
in
boxes_category
]
boxes_category
=
np
.
array
(
boxes_category
,
dtype
=
np
.
int32
)
target
.
append
(
boxes_category
)
elif
k
==
"mask"
:
elif
k
==
"mask"
:
if
"aug"
in
self
.
image_set
:
if
"aug"
in
self
.
image_set
:
mask
=
cv2
.
imread
(
self
.
masks
[
index
],
cv2
.
IMREAD_GRAYSCALE
)
mask
=
cv2
.
imread
(
self
.
masks
[
index
],
cv2
.
IMREAD_GRAYSCALE
)
...
@@ -81,9 +98,6 @@ class PascalVOC(VisionDataset):
...
@@ -81,9 +98,6 @@ class PascalVOC(VisionDataset):
mask
=
self
.
_trans_mask
(
mask
)
mask
=
self
.
_trans_mask
(
mask
)
mask
=
mask
[:,
:,
np
.
newaxis
]
mask
=
mask
[:,
:,
np
.
newaxis
]
target
.
append
(
mask
)
target
.
append
(
mask
)
elif
k
==
"boxes"
:
boxes
=
self
.
parse_voc_xml
(
ET
.
parse
(
self
.
annotations
[
index
]).
getroot
())
target
.
append
(
boxes
)
elif
k
==
"info"
:
elif
k
==
"info"
:
if
image
is
None
:
if
image
is
None
:
image
=
cv2
.
imread
(
self
.
images
[
index
],
cv2
.
IMREAD_COLOR
)
image
=
cv2
.
imread
(
self
.
images
[
index
],
cv2
.
IMREAD_COLOR
)
...
@@ -128,7 +142,6 @@ class PascalVOC(VisionDataset):
...
@@ -128,7 +142,6 @@ class PascalVOC(VisionDataset):
return
voc_dict
return
voc_dict
class_names
=
(
class_names
=
(
"background"
,
"aeroplane"
,
"aeroplane"
,
"bicycle"
,
"bicycle"
,
"bird"
,
"bird"
,
...
@@ -151,7 +164,6 @@ class PascalVOC(VisionDataset):
...
@@ -151,7 +164,6 @@ class PascalVOC(VisionDataset):
"tvmonitor"
,
"tvmonitor"
,
)
)
class_colors
=
[
class_colors
=
[
[
0
,
0
,
0
],
[
0
,
0
,
128
],
[
0
,
0
,
128
],
[
0
,
128
,
0
],
[
0
,
128
,
0
],
[
0
,
128
,
128
],
[
0
,
128
,
128
],
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录