check_dataset.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import os.path as osp
  13. import numpy as np
  14. from PIL import Image, ImageOps
  15. from .utils.visualizer import visualize
  16. from .....utils.errors import DatasetFileNotFoundError
  17. from .....utils.file_interface import custom_open
  18. from .....utils.logging import info
  19. def check_dataset(dataset_dir, output, sample_num=10):
  20. """ check dataset """
  21. dataset_dir = osp.abspath(dataset_dir)
  22. if not osp.exists(dataset_dir) or not osp.isdir(dataset_dir):
  23. raise DatasetFileNotFoundError(file_path=dataset_dir)
  24. vis_save_dir = osp.join(output, 'demo_img')
  25. if not osp.exists(vis_save_dir):
  26. os.makedirs(vis_save_dir)
  27. split_tags = ["train", "val"]
  28. attrs = dict()
  29. class_ids = set()
  30. for tag in split_tags:
  31. mapping_file = osp.join(dataset_dir, f"{tag}.txt")
  32. if not osp.exists(mapping_file):
  33. info(f"The mapping file ({mapping_file}) doesn't exist, ignored.")
  34. continue
  35. with custom_open(mapping_file, "r") as fp:
  36. lines = filter(None, (line.strip() for line in fp.readlines()))
  37. for i, line in enumerate(lines):
  38. img_file, ann_file = line.split(" ")
  39. img_file = osp.join(dataset_dir, img_file)
  40. ann_file = osp.join(dataset_dir, ann_file)
  41. assert osp.exists(img_file), FileNotFoundError(
  42. f"{img_file} not exist, please check!")
  43. assert osp.exists(ann_file), FileNotFoundError(
  44. f"{ann_file} not exist, please check!")
  45. img = np.array(
  46. ImageOps.exif_transpose(Image.open(img_file)), "uint8")
  47. ann = np.array(
  48. ImageOps.exif_transpose(Image.open(ann_file)), "uint8")
  49. assert img.shape[:2] == ann.shape, ValueError(
  50. f"The shape of {img_file}:{img.shape[:2]} and "
  51. f"{ann_file}:{ann.shape} must be the same!")
  52. class_ids = class_ids | set(ann.reshape([-1]).tolist())
  53. if i < sample_num:
  54. vis_img = visualize(img, ann)
  55. vis_img = Image.fromarray(vis_img)
  56. vis_save_path = osp.join(vis_save_dir,
  57. osp.basename(img_file))
  58. vis_img.save(vis_save_path)
  59. vis_save_path = osp.join(
  60. 'check_dataset', os.path.relpath(vis_save_path, output))
  61. if f"{tag}_sample_paths" not in attrs:
  62. attrs[f"{tag}_sample_paths"] = [vis_save_path]
  63. else:
  64. attrs[f"{tag}_sample_paths"].append(vis_save_path)
  65. if f"{tag}_samples" not in attrs:
  66. attrs[f"{tag}_samples"] = i + 1
  67. if 255 in class_ids:
  68. class_ids.remove(255)
  69. attrs["num_classes"] = len(class_ids)
  70. return attrs