未验证 提交 cf38b16f 编写于 作者: T Tingquan Gao 提交者: GitHub

Enhance dataloader robustness (#615)

* Enhance dataloader robustness

* Fix the bug about unexpected exiting when reading image failed
上级 045282fa
...@@ -70,6 +70,11 @@ int main(int argc, char **argv) { ...@@ -70,6 +70,11 @@ int main(int argc, char **argv) {
for (int idx = 0; idx < img_files_list.size(); ++idx) { for (int idx = 0; idx < img_files_list.size(); ++idx) {
std::string img_path = img_files_list[idx]; std::string img_path = img_files_list[idx];
cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR); cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR);
if (!srcimg.data) {
std::cerr << "[ERROR] image read failed! image path: " << img_path
<< "\n";
exit(-1);
}
cv::cvtColor(srcimg, srcimg, cv::COLOR_BGR2RGB); cv::cvtColor(srcimg, srcimg, cv::COLOR_BGR2RGB);
double run_time = classifier.Run(srcimg); double run_time = classifier.Run(srcimg);
......
...@@ -12,6 +12,7 @@ ...@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import random
import numpy as np import numpy as np
import imghdr import imghdr
import os import os
...@@ -182,11 +183,15 @@ class CommonDataset(Dataset): ...@@ -182,11 +183,15 @@ class CommonDataset(Dataset):
return return
def __getitem__(self, idx): def __getitem__(self, idx):
line = self.full_lines[idx] try:
img_path, label = line.split(self.delimiter) line = self.full_lines[idx]
img_path = os.path.join(self.params['data_dir'], img_path) img_path, label = line.split(self.delimiter)
with open(img_path, 'rb') as f: img_path = os.path.join(self.params['data_dir'], img_path)
img = f.read() with open(img_path, 'rb') as f:
img = f.read()
except Exception as e:
logger.error('data read faild: "{}"'.format(line))
return self.__getitem__(random.randint(0, len(self)))
return (transform(img, self.ops), int(label)) return (transform(img, self.ops), int(label))
def __len__(self): def __len__(self):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册