change behavior of -w option

object_detection/centernet/centernet.py

@@ -55,9 +55,12 @@
 # ======================
 parser = get_base_parser('CenterNet model', IMAGE_PATH, SAVE_IMAGE_PATH)
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 parser.add_argument(
     '-o', '--opset', metavar='OPSET',
@@ -186,9 +189,10 @@ def recognize_from_image(filename, detector):
     cv2.imwrite(savepath, im2show)
 
     # write prediction
-    if args.write_json:
-        pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-        write_predictions(pred_file, ary, img, category=COCO_CATEGORY, file_type='json')
+    if args.write_prediction is not None:
+        ext = args.write_prediction
+        pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+        write_predictions(pred_file, ary, img, category=COCO_CATEGORY, file_type=ext)
 
     if args.profile:
         print(detector.get_summary())
@@ -208,7 +212,7 @@ def recognize_from_video(video, detector):
     else:
         writer = None
 
-    if args.write_prediction:
+    if args.write_prediction is not None:
         frame_count = 0
         frame_digit = int(math.log10(capture.get(cv2.CAP_PROP_FRAME_COUNT)) + 1)
         video_name = os.path.splitext(os.path.basename(args.video))[0]
@@ -241,10 +245,11 @@ def recognize_from_video(video, detector):
             writer.write(img)
 
         # write prediction
-        if args.write_prediction:
+        if args.write_prediction is not None:
             savepath = get_savepath(args.savepath, video_name, post_fix = '_%s' % (str(frame_count).zfill(frame_digit) + '_res'), ext='.png')
-            pred_file = '%s.txt' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, ary, img, COCO_CATEGORY)
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, ary, img, category=COCO_CATEGORY, file_type=ext)
             frame_count += 1
 
     capture.release()

object_detection/efficientdet/efficientdet.py

@@ -94,9 +94,12 @@
     help='execute onnxruntime version.'
 )
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -298,9 +301,10 @@ def recognize_from_image(image_path, net):
     cv2.imwrite(savepath, img)
 
     # write prediction
-    if args.write_json:
-        pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-        write_predictions(pred_file, detect_object, img, category=obj_list, file_type='json')
+    if args.write_prediction is not None:
+        ext = args.write_prediction
+        pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+        write_predictions(pred_file, detect_object, img, category=obj_list, file_type=ext)
 
     if args.profile:
         print(net.get_summary())

object_detection/fastest-det/fastest-det.py

@@ -70,9 +70,12 @@
     help='IOU threshold for NMS'
 )
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -221,9 +224,10 @@ def recognize_from_image(net):
         cv2.imwrite(savepath, res_img)
 
         # write prediction
-        if args.write_json:
-            pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, detect_objects, img, category=COCO_CATEGORY, file_type='json')
+        if args.write_prediction is not None:
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, detect_objects, img, category=COCO_CATEGORY, file_type=ext)
 
     logger.info('Script finished successfully.')
 

object_detection/m2det/m2det.py

@@ -56,9 +56,12 @@
 # ======================
 parser = get_base_parser('m2det model', IMAGE_PATH, SAVE_IMAGE_PATH)
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -223,9 +226,10 @@ def recognize_from_image(filename, detector):
     cv2.imwrite(savepath, img)
 
     # write prediction
-    if args.write_json:
-        pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-        write_predictions(pred_file, detect_object, img, category=COCO_CATEGORY, file_type='json')
+    if args.write_prediction is not None:
+        ext = args.write_prediction
+        pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+        write_predictions(pred_file, detect_object, img, category=COCO_CATEGORY, file_type=ext)
 
     if args.profile:
         print(detector.get_summary())

object_detection/mobilenet_ssd/mobilenet_ssd.py

@@ -38,9 +38,12 @@
     help='model lists: ' + ' | '.join(MODEL_LISTS) + ' (default: mb2-ssd-lite)'
 )
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -128,9 +131,10 @@ def recognize_from_image():
         cv2.imwrite(savepath, res_img)
 
         # write prediction
-        if args.write_json:
-            pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, detector, org_img, category=VOC_CATEGORY, file_type='json')
+        if args.write_prediction is not None:
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, detector, org_img, category=VOC_CATEGORY, file_type=ext)
 
     if args.profile:
         print(detector.get_summary())

object_detection/nanodet/nanodet.py

@@ -67,9 +67,12 @@
          'nanodet_m, nanodet_m_416, nanodet_t, nanodet-RepVGG-A0_416]'
 )
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -137,9 +140,10 @@ def recognize_from_image():
         cv2.imwrite(savepath, res_img)
 
         # write prediction
-        if args.write_json:
-            pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, detect_object, raw_img, category=COCO_CATEGORY, file_type='json')
+        if args.write_prediction is not None:
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, detect_object, raw_img, category=COCO_CATEGORY, file_type=ext)
 
     logger.info('Script finished successfully.')
 

object_detection/pedestrian_detection/pedestrian_detection.py

@@ -60,9 +60,12 @@
     help='The detection height and height for yolo. (default: 416)'
 )
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -117,9 +120,10 @@ def recognize_from_image():
         cv2.imwrite(savepath, res_img)
 
         # write prediction
-        if args.write_json:
-            pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, detector, img, category=COCO_CATEGORY, file_type='json')
+        if args.write_prediction is not None:
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, detector, img, category=COCO_CATEGORY, file_type=ext)
 
     if args.profile:
         print(detector.get_summary())

object_detection/picodet/picodet.py

@@ -87,9 +87,12 @@
     help='model type'
 )
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -232,9 +235,10 @@ def recognize_from_image(net):
         cv2.imwrite(savepath, res_img)
 
         # write prediction
-        if args.write_json:
-            pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, detect_object, img, category=COCO_CATEGORY, file_type='json')
+        if args.write_prediction is not None:
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, detect_object, img, category=COCO_CATEGORY, file_type=ext)
 
     logger.info('Script finished successfully.')
 

object_detection/yolor/yolor.py

@@ -36,9 +36,12 @@
     help='yolor_p6, yolor_w6]'
 )
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 REMOTE_PATH = 'https://storage.googleapis.com/ailia-models/yolor/'
@@ -109,9 +112,10 @@ def recognize_from_image():
         cv2.imwrite(savepath, res_img)
 
         # write prediction
-        if args.write_json:
-            pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, detect_object, raw_img, category=COCO_CATEGORY, file_type='json')
+        if args.write_prediction is not None:
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, detect_object, raw_img, category=COCO_CATEGORY, file_type=ext)
 
     logger.info('Script finished successfully.')
 

object_detection/yolov1-tiny/yolov1-tiny.py

@@ -60,9 +60,12 @@
 # ======================
 parser = get_base_parser('Yolov1 tiny model', IMAGE_PATH, SAVE_IMAGE_PATH)
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -111,9 +114,10 @@ def recognize_from_image():
         cv2.imwrite(savepath, res_img)
 
         # write prediction
-        if args.write_json:
-            pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, detector, img, category=VOC_CATEGORY, file_type='json')
+        if args.write_prediction is not None:
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, detector, img, category=VOC_CATEGORY, file_type=ext)
 
     if args.profile:
         print(detector.get_summary())

object_detection/yolov2-tiny/yolov2-tiny.py

@@ -86,9 +86,12 @@
     metavar='DATASET', default='coco'
 )
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -172,9 +175,10 @@ def recognize_from_image():
             cv2.imwrite(savepath, res_img)
 
             # write prediction
-            if args.write_json:
-                pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-                write_predictions(pred_file, detector, img, category=CATEGORY, file_type='json')
+            if args.write_prediction is not None:
+                ext = args.write_prediction
+                pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+                write_predictions(pred_file, detector, img, category=CATEGORY, file_type=ext)
 
             if args.profile:
                 print(detector.get_summary())
@@ -188,10 +192,11 @@ def recognize_from_image():
                 logger.info(f'saved at : {savepath}')
                 cv2.imwrite(savepath, res_img)
 
-                 # write prediction
-                if args.write_json:
-                    pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-                    write_predictions(pred_file, detector, img, category=CATEGORY, file_type='json')
+                # write prediction
+                if args.write_prediction is not None:
+                    ext = args.write_prediction
+                    pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+                    write_predictions(pred_file, detector, img, category=CATEGORY, file_type=ext)
 
                 if args.profile:
                     print(detector.get_summary())

object_detection/yolov2/yolov2.py

@@ -59,9 +59,12 @@
 # ======================
 parser = get_base_parser('Yolov2 model', IMAGE_PATH, SAVE_IMAGE_PATH)
 parser.add_argument(
-    '-w', '--write_json',
-    action='store_true',
-    help='Flag to output results to json file.'
+    '-w', '--write_prediction',
+    nargs='?',
+    const='txt',
+    choices=['txt', 'json'],
+    type=str,
+    help='Output results to txt or json file.'
 )
 args = update_parser(parser)
 
@@ -111,9 +114,10 @@ def recognize_from_image():
         cv2.imwrite(savepath, res_img)
 
         # write prediction
-        if args.write_json:
-            pred_file = '%s.json' % savepath.rsplit('.', 1)[0]
-            write_predictions(pred_file, detector, img, category=COCO_CATEGORY, file_type='json')
+        if args.write_prediction is not None:
+            ext = args.write_prediction
+            pred_file = "%s.%s" % (savepath.rsplit('.', 1)[0], ext)
+            write_predictions(pred_file, detector, img, category=COCO_CATEGORY, file_type=ext)
 
     if args.profile:
         print(detector.get_summary())