Browse Source

Create on_docker.py

LaraStuStu 5 years ago
parent
commit
af74bc3a72
1 changed files with 100 additions and 0 deletions
  1. 100 0
      DataAnnotation/labelme/labelme/cli/on_docker.py

+ 100 - 0
DataAnnotation/labelme/labelme/cli/on_docker.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python
+
+from __future__ import print_function
+
+import argparse
+import distutils.spawn
+import json
+import os
+import os.path as osp
+import platform
+import shlex
+import subprocess
+import sys
+
+
+def get_ip():
+    dist = platform.platform().split('-')[0]
+    if dist == 'Linux':
+        return ''
+    elif dist == 'Darwin':
+        cmd = 'ifconfig en0'
+        output = subprocess.check_output(shlex.split(cmd))
+        if str != bytes:  # Python3
+            output = output.decode('utf-8')
+        for row in output.splitlines():
+            cols = row.strip().split(' ')
+            if cols[0] == 'inet':
+                ip = cols[1]
+                return ip
+        else:
+            raise RuntimeError('No ip is found.')
+    else:
+        raise RuntimeError('Unsupported platform.')
+
+
+def labelme_on_docker(in_file, out_file):
+    ip = get_ip()
+    cmd = 'xhost + %s' % ip
+    subprocess.check_output(shlex.split(cmd))
+
+    if out_file:
+        out_file = osp.abspath(out_file)
+        if osp.exists(out_file):
+            raise RuntimeError('File exists: %s' % out_file)
+        else:
+            open(osp.abspath(out_file), 'w')
+
+    cmd = 'docker run -it --rm' \
+        ' -e DISPLAY={0}:0' \
+        ' -e QT_X11_NO_MITSHM=1' \
+        ' -v /tmp/.X11-unix:/tmp/.X11-unix' \
+        ' -v {1}:{2}' \
+        ' -w /home/developer'
+    in_file_a = osp.abspath(in_file)
+    in_file_b = osp.join('/home/developer', osp.basename(in_file))
+    cmd = cmd.format(
+        ip,
+        in_file_a,
+        in_file_b,
+    )
+    if out_file:
+        out_file_a = osp.abspath(out_file)
+        out_file_b = osp.join('/home/developer', osp.basename(out_file))
+        cmd += ' -v {0}:{1}'.format(out_file_a, out_file_b)
+    cmd += ' wkentaro/labelme labelme {0}'.format(in_file_b)
+    if out_file:
+        cmd += ' -O {0}'.format(out_file_b)
+    subprocess.call(shlex.split(cmd))
+
+    if out_file:
+        try:
+            json.load(open(out_file))
+            return out_file
+        except Exception:
+            if open(out_file).read() == '':
+                os.remove(out_file)
+            raise RuntimeError('Annotation is cancelled.')
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('in_file', help='Input file or directory.')
+    parser.add_argument('-O', '--output')
+    args = parser.parse_args()
+
+    if not distutils.spawn.find_executable('docker'):
+        print('Please install docker', file=sys.stderr)
+        sys.exit(1)
+
+    try:
+        out_file = labelme_on_docker(args.in_file, args.output)
+        if out_file:
+            print('Saved to: %s' % out_file)
+    except RuntimeError as e:
+        sys.stderr.write(e.__str__() + '\n')
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    main()