segment.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. # @Author : lightXu
  2. # @File : segment.py
  3. import os
  4. import xml.etree.cElementTree as ET
  5. from segment.image_operation import utils
  6. def joint_image(raw_img_path, bbox, lines_list):
  7. lines_dir = raw_img_path.replace('.jpg', '_lines')
  8. lines_file_list = os.listdir(lines_dir)
  9. lines_file_list = sorted([ele.replace('jpg', '')
  10. for ele in lines_file_list if ele.endswith('.jpg')])
  11. exam_items_bbox = []
  12. tree = ET.parse(r'./segment/exam_info/000000-template.xml') # xml tree
  13. for index_num, j in enumerate(lines_list):
  14. if j[1] == j[0]:
  15. continue
  16. elif j[1] - j[0] == 1:
  17. index_list = lines_file_list[j[0]].split('_')
  18. y_low = int(index_list[0])
  19. y_high = int(index_list[1])
  20. x_low = int(index_list[2])
  21. x_high = int(index_list[3])
  22. else:
  23. index_list0 = lines_file_list[j[0]].split('_') # [33, 37]
  24. index_list1 = lines_file_list[j[1] - 1].split('_')
  25. y_low = int(index_list0[0])
  26. y_high = int(index_list1[1])
  27. tmp_x_low_list = [ele.split('_')[2]
  28. for ele in lines_file_list[j[0]:j[1]]]
  29. tmp_x_high_list = [ele.split('_')[3]
  30. for ele in lines_file_list[j[0]:j[1]]]
  31. x_low = int(min(tmp_x_low_list))
  32. x_high = int(max(tmp_x_high_list))
  33. exam_bbox = [bbox[2] + x_low, bbox[0] + y_low, bbox[2] + x_high, bbox[0] + y_high]
  34. tree = utils.create_xml('{:02d}'.format(index_num), tree,
  35. exam_bbox[0], exam_bbox[1], exam_bbox[2], exam_bbox[3])
  36. exam_items_bbox.append(exam_bbox)
  37. # print(exam_items_bbox)
  38. tree.write(raw_img_path.replace('.jpg', '.xml'))
  39. return exam_items_bbox