@@ -280,15 +280,15 @@ def predict(
280
280
postprocess_match_metric : str = "IOS" ,
281
281
postprocess_match_threshold : float = 0.5 ,
282
282
postprocess_class_agnostic : bool = False ,
283
- export_visual : bool = True ,
283
+ export_visual : bool = False ,
284
284
export_pickle : bool = False ,
285
285
export_crop : bool = False ,
286
286
dataset_json_path : bool = None ,
287
287
project : str = "runs/predict" ,
288
288
name : str = "exp" ,
289
- visual_bbox_thickness : int = 1 ,
290
- visual_text_size : float = 0.3 ,
291
- visual_text_thickness : int = 1 ,
289
+ visual_bbox_thickness : int = None ,
290
+ visual_text_size : float = None ,
291
+ visual_text_thickness : int = None ,
292
292
visual_export_format : str = "png" ,
293
293
verbose : int = 1 ,
294
294
):
@@ -469,43 +469,44 @@ def predict(
469
469
coco_prediction_json = coco_prediction .json
470
470
if coco_prediction_json ["bbox" ]:
471
471
coco_json .append (coco_prediction_json )
472
- # convert ground truth annotations to object_prediction_list
473
- coco_image : CocoImage = coco .images [ind ]
474
- object_prediction_gt_list : List [ObjectPrediction ] = []
475
- for coco_annotation in coco_image .annotations :
476
- coco_annotation_dict = coco_annotation .json
477
- category_name = coco_annotation .category_name
478
- full_shape = [coco_image .height , coco_image .width ]
479
- object_prediction_gt = ObjectPrediction .from_coco_annotation_dict (
480
- annotation_dict = coco_annotation_dict , category_name = category_name , full_shape = full_shape
472
+ if export_visual :
473
+ # convert ground truth annotations to object_prediction_list
474
+ coco_image : CocoImage = coco .images [ind ]
475
+ object_prediction_gt_list : List [ObjectPrediction ] = []
476
+ for coco_annotation in coco_image .annotations :
477
+ coco_annotation_dict = coco_annotation .json
478
+ category_name = coco_annotation .category_name
479
+ full_shape = [coco_image .height , coco_image .width ]
480
+ object_prediction_gt = ObjectPrediction .from_coco_annotation_dict (
481
+ annotation_dict = coco_annotation_dict , category_name = category_name , full_shape = full_shape
482
+ )
483
+ object_prediction_gt_list .append (object_prediction_gt )
484
+ # export visualizations with ground truths
485
+ output_dir = str (visual_with_gt_dir / Path (relative_filepath ).parent )
486
+ color = (0 , 255 , 0 ) # original annotations in green
487
+ result = visualize_object_predictions (
488
+ np .ascontiguousarray (image_as_pil ),
489
+ object_prediction_list = object_prediction_gt_list ,
490
+ rect_th = visual_bbox_thickness ,
491
+ text_size = visual_text_size ,
492
+ text_th = visual_text_thickness ,
493
+ color = color ,
494
+ output_dir = None ,
495
+ file_name = None ,
496
+ export_format = None ,
497
+ )
498
+ color = (255 , 0 , 0 ) # model predictions in red
499
+ _ = visualize_object_predictions (
500
+ result ["image" ],
501
+ object_prediction_list = object_prediction_list ,
502
+ rect_th = visual_bbox_thickness ,
503
+ text_size = visual_text_size ,
504
+ text_th = visual_text_thickness ,
505
+ color = color ,
506
+ output_dir = output_dir ,
507
+ file_name = filename_without_extension ,
508
+ export_format = visual_export_format ,
481
509
)
482
- object_prediction_gt_list .append (object_prediction_gt )
483
- # export visualizations with ground truths
484
- output_dir = str (visual_with_gt_dir / Path (relative_filepath ).parent )
485
- color = (0 , 255 , 0 ) # original annotations in green
486
- result = visualize_object_predictions (
487
- np .ascontiguousarray (image_as_pil ),
488
- object_prediction_list = object_prediction_gt_list ,
489
- rect_th = visual_bbox_thickness ,
490
- text_size = visual_text_size ,
491
- text_th = visual_text_thickness ,
492
- color = color ,
493
- output_dir = None ,
494
- file_name = None ,
495
- export_format = None ,
496
- )
497
- color = (255 , 0 , 0 ) # model predictions in red
498
- _ = visualize_object_predictions (
499
- result ["image" ],
500
- object_prediction_list = object_prediction_list ,
501
- rect_th = visual_bbox_thickness ,
502
- text_size = visual_text_size ,
503
- text_th = visual_text_thickness ,
504
- color = color ,
505
- output_dir = output_dir ,
506
- file_name = filename_without_extension ,
507
- export_format = visual_export_format ,
508
- )
509
510
510
511
time_start = time .time ()
511
512
# export prediction boxes
0 commit comments