vehicle_dataset.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import print_function
  15. import numpy as np
  16. import paddle
  17. from paddle.io import Dataset
  18. import os
  19. import cv2
  20. from paddlex.ppcls.data import preprocess
  21. from paddlex.ppcls.data.preprocess import transform
  22. from paddlex.ppcls.utils import logger
  23. from .common_dataset import create_operators
  24. class CompCars(Dataset):
  25. def __init__(self,
  26. image_root,
  27. cls_label_path,
  28. label_root=None,
  29. transform_ops=None,
  30. bbox_crop=False):
  31. self._img_root = image_root
  32. self._cls_path = cls_label_path
  33. self._label_root = label_root
  34. if transform_ops:
  35. self._transform_ops = create_operators(transform_ops)
  36. self._bbox_crop = bbox_crop
  37. self._dtype = paddle.get_default_dtype()
  38. self._load_anno()
  39. def _load_anno(self):
  40. assert os.path.exists(self._cls_path)
  41. assert os.path.exists(self._img_root)
  42. if self._bbox_crop:
  43. assert os.path.exists(self._label_root)
  44. self.images = []
  45. self.labels = []
  46. self.bboxes = []
  47. with open(self._cls_path) as fd:
  48. lines = fd.readlines()
  49. for l in lines:
  50. l = l.strip().split()
  51. if not self._bbox_crop:
  52. self.images.append(os.path.join(self._img_root, l[0]))
  53. self.labels.append(int(l[1]))
  54. else:
  55. label_path = os.path.join(self._label_root,
  56. l[0].split('.')[0] + '.txt')
  57. assert os.path.exists(label_path)
  58. with open(label_path) as f:
  59. bbox = f.readlines()[-1].strip().split()
  60. bbox = [int(x) for x in bbox]
  61. self.images.append(os.path.join(self._img_root, l[0]))
  62. self.labels.append(int(l[1]))
  63. self.bboxes.append(bbox)
  64. assert os.path.exists(self.images[-1])
  65. def __getitem__(self, idx):
  66. img = cv2.imread(self.images[idx])
  67. img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  68. if self._bbox_crop:
  69. bbox = self.bboxes[idx]
  70. img = img[bbox[1]:bbox[3], bbox[0]:bbox[2], :]
  71. if self._transform_ops:
  72. img = transform(img, self._transform_ops)
  73. img = img.transpose((2, 0, 1))
  74. return (img, self.labels[idx])
  75. def __len__(self):
  76. return len(self.images)
  77. @property
  78. def class_num(self):
  79. return len(set(self.labels))
  80. class VeriWild(Dataset):
  81. def __init__(
  82. self,
  83. image_root,
  84. cls_label_path,
  85. transform_ops=None, ):
  86. self._img_root = image_root
  87. self._cls_path = cls_label_path
  88. if transform_ops:
  89. self._transform_ops = create_operators(transform_ops)
  90. self._dtype = paddle.get_default_dtype()
  91. self._load_anno()
  92. def _load_anno(self):
  93. assert os.path.exists(self._cls_path)
  94. assert os.path.exists(self._img_root)
  95. self.images = []
  96. self.labels = []
  97. self.cameras = []
  98. with open(self._cls_path) as fd:
  99. lines = fd.readlines()
  100. for l in lines:
  101. l = l.strip().split()
  102. self.images.append(os.path.join(self._img_root, l[0]))
  103. self.labels.append(np.int64(l[1]))
  104. self.cameras.append(np.int64(l[2]))
  105. assert os.path.exists(self.images[-1])
  106. def __getitem__(self, idx):
  107. try:
  108. with open(self.images[idx], 'rb') as f:
  109. img = f.read()
  110. if self._transform_ops:
  111. img = transform(img, self._transform_ops)
  112. img = img.transpose((2, 0, 1))
  113. return (img, self.labels[idx], self.cameras[idx])
  114. except Exception as ex:
  115. logger.error("Exception occured when parse line: {} with msg: {}".
  116. format(self.images[idx], ex))
  117. rnd_idx = np.random.randint(self.__len__())
  118. return self.__getitem__(rnd_idx)
  119. def __len__(self):
  120. return len(self.images)
  121. @property
  122. def class_num(self):
  123. return len(set(self.labels))