Sort images and check the input image is assumed image format. (#246)

segmentation/image_demo.py

@@ -15,6 +15,12 @@ import os
 
 
 def test_single_image(model, img_name, out_dir, color_palette, opacity):
+    # check img_name is an image file or not
+    assumed_imgformat = ('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')
+    if (not img_name.lower().endswith(assumed_imgformat)):
+        print(f"Skip {img_name} because it is not an image file.")
+        return
+
     result = inference_segmentor(model, img_name)
 
     # show the results
@@ -33,7 +39,8 @@ def test_single_image(model, img_name, out_dir, color_palette, opacity):
 
 def main():
     parser = ArgumentParser()
-    parser.add_argument('img', help='Image file or a directory contains images')
+    parser.add_argument(
+        'img', help='Image file or a directory contains images')
     parser.add_argument('config', help='Config file')
     parser.add_argument('checkpoint', help='Checkpoint file')
     parser.add_argument('--out', type=str, default="demo", help='out dir')
@@ -61,10 +68,13 @@ def main():
 
     # check arg.img is directory of a single image.
     if osp.isdir(args.img):
-        for img in os.listdir(args.img):
-            test_single_image(model, osp.join(args.img, img), args.out, get_palette(args.palette), args.opacity)
+        for img in sorted(os.listdir(args.img)):
+            test_single_image(model, osp.join(args.img, img),
+                              args.out, get_palette(args.palette), args.opacity)
     else:
-        test_single_image(model, args.img, args.out, get_palette(args.palette), args.opacity)
+        test_single_image(model, args.img, args.out,
+                          get_palette(args.palette), args.opacity)
+
 
 if __name__ == '__main__':
     main()