55import onnxruntime as ort
66from matplotlib .colors import TABLEAU_COLORS
77
8- ALLOWED_EXTENSIONS = {' txt' , ' pdf' , ' png' , ' jpg' , ' jpeg' , ' gif' }
8+ ALLOWED_EXTENSIONS = {" txt" , " pdf" , " png" , " jpg" , " jpeg" , " gif" }
99BASE_DIR = os .path .dirname (os .path .abspath (__file__ ))
1010h , w = 640 , 640
1111model_onnx_path = os .path .join (BASE_DIR , "yolov7-p6-bonefracture.onnx" )
@@ -44,7 +44,7 @@ def preproc(img):
4444
4545def model_inference (model_path , image_np , device = "cpu" ):
4646
47- providers = [' CUDAExecutionProvider' ] if device == "cuda" else [' CPUExecutionProvider' ]
47+ providers = [" CUDAExecutionProvider" ] if device == "cuda" else [" CPUExecutionProvider" ]
4848 session = ort .InferenceSession (model_path , providers = providers )
4949 input_name = session .get_inputs ()[0 ].name
5050 output_name = session .get_outputs ()[0 ].name
@@ -97,9 +97,9 @@ def post_process(img, output, score_threshold=0.3, format="xywh"):
9797
9898if __name__ == "__main__" :
9999
100- st .title (' Bone Fracture Detection' )
100+ st .title (" Bone Fracture Detection" )
101101
102- uploaded_file = st .file_uploader ("Choose a image file" , type = [' png' , ' jpg' , ' jpeg' , ' gif' ])
102+ uploaded_file = st .file_uploader ("Choose a image file" , type = [" png" , " jpg" , " jpeg" , " gif" ])
103103
104104 if uploaded_file is not None :
105105
@@ -116,7 +116,7 @@ def post_process(img, output, score_threshold=0.3, format="xywh"):
116116
117117 st .download_button (
118118 label = "Download prediction" ,
119- data = cv2 .imencode (' .png' , out_img [..., ::- 1 ])[1 ].tobytes (),
119+ data = cv2 .imencode (" .png" , out_img [..., ::- 1 ])[1 ].tobytes (),
120120 file_name = uploaded_file .name ,
121121 mime = "image/png"
122122 )
0 commit comments