check_dataset.py 3.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  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 json
  12. import os
  13. import os.path as osp
  14. from collections import defaultdict, Counter
  15. from pathlib import Path
  16. from PIL import Image, ImageOps
  17. from pycocotools.coco import COCO
  18. from .utils.visualizer import draw_bbox, draw_mask
  19. from .....utils.errors import DatasetFileNotFoundError
  20. from .....utils.logging import info
  21. def check(dataset_dir, output, sample_num=10):
  22. """ check dataset """
  23. info(dataset_dir)
  24. dataset_dir = osp.abspath(dataset_dir)
  25. if not osp.exists(dataset_dir) or not osp.isdir(dataset_dir):
  26. raise DatasetFileNotFoundError(file_path=dataset_dir)
  27. sample_cnts = dict()
  28. sample_paths = defaultdict(list)
  29. im_sizes = defaultdict(Counter)
  30. tags = ['instance_train', 'instance_val']
  31. for _, tag in enumerate(tags):
  32. file_list = osp.join(dataset_dir, f'annotations/{tag}.json')
  33. if not osp.exists(file_list):
  34. if tag in ('instance_train', 'instance_val'):
  35. # train and val file lists must exist
  36. raise DatasetFileNotFoundError(
  37. file_path=file_list,
  38. solution=f"Ensure that both `instance_train.json` and `instance_val.json` exist in \
  39. {dataset_dir}/annotations")
  40. else:
  41. continue
  42. else:
  43. with open(file_list, 'r', encoding='utf-8') as f:
  44. jsondata = json.load(f)
  45. datanno = jsondata['annotations']
  46. sample_cnts[tag] = len(datanno)
  47. coco = COCO(file_list)
  48. num_class = len(coco.getCatIds())
  49. vis_save_dir = osp.join(output, 'demo_img')
  50. image_info = jsondata['images']
  51. for i in range(sample_num):
  52. file_name = image_info[i]['file_name']
  53. img_id = image_info[i]['id']
  54. img_path = osp.join(dataset_dir, 'images', file_name)
  55. if not osp.exists(img_path):
  56. raise DatasetFileNotFoundError(file_path=img_path)
  57. img = Image.open(img_path)
  58. img = ImageOps.exif_transpose(img)
  59. vis_im = draw_bbox(img, coco, img_id)
  60. vis_im = draw_mask(vis_im, coco, img_id)
  61. vis_path = osp.join(vis_save_dir, file_name)
  62. Path(vis_path).parent.mkdir(parents=True, exist_ok=True)
  63. vis_im.save(vis_path)
  64. sample_path = osp.join('check_dataset',
  65. os.path.relpath(vis_path, output))
  66. sample_paths[tag].append(sample_path)
  67. attrs = {}
  68. attrs['num_classes'] = num_class
  69. attrs['train_samples'] = sample_cnts['instance_train']
  70. attrs['train_sample_paths'] = sample_paths['instance_train']
  71. attrs['val_samples'] = sample_cnts['instance_val']
  72. attrs['val_sample_paths'] = sample_paths['instance_val']
  73. return attrs