labelme_on_docker 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. #!/usr/bin/env python
  2. import argparse
  3. import json
  4. import os
  5. import os.path as osp
  6. import platform
  7. import shlex
  8. import subprocess
  9. import sys
  10. def get_ip():
  11. dist = platform.platform().split('-')[0]
  12. if dist == 'Linux':
  13. return ''
  14. elif dist == 'Darwin':
  15. cmd = 'ifconfig en0'
  16. output = subprocess.check_output(shlex.split(cmd))
  17. for row in output.splitlines():
  18. cols = row.strip().split(' ')
  19. if cols[0] == 'inet':
  20. ip = cols[1]
  21. return ip
  22. else:
  23. raise RuntimeError('No ip is found.')
  24. else:
  25. raise RuntimeError('Unsupported platform.')
  26. def labelme_on_docker(image_file, out_file):
  27. ip = get_ip()
  28. cmd = 'xhost + %s' % ip
  29. subprocess.check_output(shlex.split(cmd))
  30. out_file = osp.abspath(out_file)
  31. if osp.exists(out_file):
  32. raise RuntimeError('File exists: %s' % out_file)
  33. else:
  34. open(osp.abspath(out_file), 'w')
  35. cmd = 'docker run -it --rm' \
  36. ' -e DISPLAY={0}:0' \
  37. ' -e QT_X11_NO_MITSHM=1' \
  38. ' -v /tmp/.X11-unix:/tmp/.X11-unix' \
  39. ' -v {1}:{2}' \
  40. ' -v {3}:{4}' \
  41. ' -w /home/developer' \
  42. ' labelme' \
  43. ' labelme {2} -O {4}'
  44. cmd = cmd.format(
  45. ip,
  46. osp.abspath(image_file),
  47. osp.join('/home/developer', osp.basename(image_file)),
  48. osp.abspath(out_file),
  49. osp.join('/home/developer', osp.basename(out_file)),
  50. )
  51. subprocess.call(shlex.split(cmd))
  52. try:
  53. json.load(open(out_file))
  54. return out_file
  55. except Exception as e:
  56. if open(out_file).read() == '':
  57. os.remove(out_file)
  58. raise RuntimeError('Annotation is cancelled.')
  59. def main():
  60. parser = argparse.ArgumentParser()
  61. parser.add_argument('image_file')
  62. parser.add_argument('-O', '--output', required=True)
  63. args = parser.parse_args()
  64. try:
  65. out_file = labelme_on_docker(args.image_file, args.output)
  66. print('Saved to: %s' % out_file)
  67. except RuntimeError as e:
  68. sys.stderr.write(e.__str__() + '\n')
  69. sys.exit(1)
  70. if __name__ == '__main__':
  71. main()