example-apps/PDD/pcb-defect-detection/data/io/divide_data.py

   1 # -*- coding: utf-8 -*-
   2 from __future__ import absolute_import, division, print_function
   3
   4 import math
   5 import os
   6 import random
   7 import shutil
   8 import sys
   9
  10 sys.path.append('../../')
  11
  12 from libs.configs import cfgs
  13
  14
  15 def mkdir(path):
  16   if not os.path.exists(path):
  17     os.makedirs(path)
  18
  19
  20 divide_rate = 0.8
  21
  22 image_path = os.path.join(cfgs.ROOT_PATH, '{}/JPEGImages'.format(cfgs.DATASET_NAME))
  23 xml_path = os.path.join(cfgs.ROOT_PATH, '{}/Annotations'.format(cfgs.DATASET_NAME))
  24
  25 image_list = os.listdir(image_path)
  26
  27 image_name = [n.split('.')[0] for n in image_list]
  28
  29 random.shuffle(image_name)
  30
  31 train_image = image_name[:int(math.ceil(len(image_name)) * divide_rate)]
  32 test_image = image_name[int(math.ceil(len(image_name)) * divide_rate):]
  33
  34 image_output_train = os.path.join(
  35     cfgs.ROOT_PATH, '{}_train/JPEGImages'.format(cfgs.DATASET_NAME))
  36 mkdir(image_output_train)
  37 image_output_test = os.path.join(
  38     cfgs.ROOT_PATH, '{}_test/JPEGImages'.format(cfgs.DATASET_NAME))
  39 mkdir(image_output_test)
  40
  41 xml_train = os.path.join(cfgs.ROOT_PATH, '{}_train/Annotations'.format(cfgs.DATASET_NAME))
  42 mkdir(xml_train)
  43 xml_test = os.path.join(cfgs.ROOT_PATH, '{}_test/Annotations'.format(cfgs.DATASET_NAME))
  44 mkdir(xml_test)
  45
  46
  47 count = 0
  48 for i in train_image:
  49   shutil.copy(os.path.join(image_path, i + '.jpg'), image_output_train)
  50   if os.path.exists(os.path.join(xml_path, i + '.xml')):
  51     shutil.copy(os.path.join(xml_path, i + '.xml'), xml_train)
  52   if count % 1000 == 0:
  53     print("process step {}".format(count))
  54   count += 1
  55
  56 for i in test_image:
  57   shutil.copy(os.path.join(image_path, i + '.jpg'), image_output_test)
  58   shutil.copy(os.path.join(xml_path, i + '.xml'), xml_test)
  59   if count % 1000 == 0:
  60     print("process step {}".format(count))
  61   count += 1