This repository was archived by the owner on Dec 2, 2022. It is now read-only.
File tree Expand file tree Collapse file tree 2 files changed +28
-21
lines changed Expand file tree Collapse file tree 2 files changed +28
-21
lines changed Original file line number Diff line number Diff line change @@ -60,6 +60,7 @@ from craft_text_detector import (
6060 get_prediction,
6161 export_detected_regions,
6262 export_extra_results,
63+ empty_cuda_cache
6364)
6465
6566# set image path and export folder directory
@@ -75,31 +76,34 @@ craft_net = load_craftnet_model(cuda=True)
7576
7677# perform prediction
7778prediction_result = get_prediction(
78- image = image,
79- craft_net = craft_net,
80- refine_net = refine_net,
81- text_threshold = 0.7 ,
82- link_threshold = 0.4 ,
83- low_text = 0.4 ,
84- cuda = True ,
85- long_size = 1280
79+ image = image,
80+ craft_net = craft_net,
81+ refine_net = refine_net,
82+ text_threshold = 0.7 ,
83+ link_threshold = 0.4 ,
84+ low_text = 0.4 ,
85+ cuda = True ,
86+ long_size = 1280
8687)
8788
8889# export detected text regions
8990exported_file_paths = export_detected_regions(
90- image_path = image_path,
91- image = image,
92- regions = prediction_result[" boxes" ],
93- output_dir = output_dir,
94- rectify = True
91+ image_path = image_path,
92+ image = image,
93+ regions = prediction_result[" boxes" ],
94+ output_dir = output_dir,
95+ rectify = True
9596)
9697
9798# export heatmap, detection points, box visualization
9899export_extra_results(
99- image_path = image_path,
100- image = image,
101- regions = prediction_result[" boxes" ],
102- heatmaps = prediction_result[" heatmaps" ],
103- output_dir = output_dir
100+ image_path = image_path,
101+ image = image,
102+ regions = prediction_result[" boxes" ],
103+ heatmaps = prediction_result[" heatmaps" ],
104+ output_dir = output_dir
104105)
106+
107+ # unload models from gpu
108+ empty_cuda_cache()
105109```
Original file line number Diff line number Diff line change 44import craft_text_detector .file_utils as file_utils
55import craft_text_detector .image_utils as image_utils
66import craft_text_detector .predict as predict
7+ import craft_text_detector .torch_utils as torch_utils
78
8- __version__ = "0.3.0 "
9+ __version__ = "0.3.1 "
910
1011
1112__all__ = [
1516 "get_prediction" ,
1617 "export_detected_regions" ,
1718 "export_extra_results" ,
19+ "empty_cuda_cache" ,
1820 "Craft" ,
1921]
2022
2426get_prediction = predict .get_prediction
2527export_detected_regions = file_utils .export_detected_regions
2628export_extra_results = file_utils .export_extra_results
29+ empty_cuda_cache = torch_utils .empty_cuda_cache
2730
2831
2932class Craft :
@@ -90,14 +93,14 @@ def unload_craftnet_model(self):
9093 Unloads craftnet model
9194 """
9295 self .craft_net = None
93- craft_utils . empty_cuda_cache ()
96+ empty_cuda_cache ()
9497
9598 def unload_refinenet_model (self ):
9699 """
97100 Unloads refinenet model
98101 """
99102 self .refine_net = None
100- craft_utils . empty_cuda_cache ()
103+ empty_cuda_cache ()
101104
102105 def detect_text (self , image_path ):
103106 """
You can’t perform that action at this time.
0 commit comments