Browse Source

Update labelme_on_docker to handle input directory

Kentaro Wada 7 năm trước cách đây
mục cha
commit
7c3e465de3
1 tập tin đã thay đổi với 32 bổ sung25 xóa
  1. 32 25
      labelme/cli/on_docker.py

+ 32 - 25
labelme/cli/on_docker.py

@@ -30,53 +30,60 @@ def get_ip():
         raise RuntimeError('Unsupported platform.')
 
 
-def labelme_on_docker(image_file, out_file):
+def labelme_on_docker(in_file, out_file):
     ip = get_ip()
     cmd = 'xhost + %s' % ip
     subprocess.check_output(shlex.split(cmd))
 
-    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')
+    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}' \
-        ' -v {3}:{4}' \
-        ' -w /home/developer' \
-        ' wkentaro/labelme' \
-        ' labelme {2} -O {4}'
+        ' -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,
-        osp.abspath(image_file),
-        osp.join('/home/developer', osp.basename(image_file)),
-        osp.abspath(out_file),
-        osp.join('/home/developer', osp.basename(out_file)),
+        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))
 
-    try:
-        json.load(open(out_file))
-        return out_file
-    except Exception as e:
-        if open(out_file).read() == '':
-            os.remove(out_file)
-        raise RuntimeError('Annotation is cancelled.')
+    if out_file:
+        try:
+            json.load(open(out_file))
+            return out_file
+        except Exception as e:
+            if open(out_file).read() == '':
+                os.remove(out_file)
+            raise RuntimeError('Annotation is cancelled.')
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('image_file')
-    parser.add_argument('-O', '--output', required=True)
+    parser.add_argument('in_file', help='Input file or directory.')
+    parser.add_argument('-O', '--output')
     args = parser.parse_args()
 
     try:
-        out_file = labelme_on_docker(args.image_file, args.output)
-        print('Saved to: %s' % out_file)
+        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)