set intra op threads (#5)

- set intra op threads (7e156dc8b9521e9a3d42abe2727a9272891ca273)


Co-authored-by: Sayak Paul <[email protected]>

app.py

@@ -1,31 +1,36 @@
-import gradio as gr
-
-import sys
 import csv
-import numpy as np
+import os
+import sys
+
 import cv2
-from matplotlib import gridspec
+import gradio as gr
 import matplotlib.pyplot as plt
+import numpy as np
 import onnxruntime as ort
+from matplotlib import gridspec
 
 ade_palette = []
 labels_list = []
 
 csv.field_size_limit(sys.maxsize)
 
-with open(r'labels.txt', 'r') as fp:
+with open(r"labels.txt", "r") as fp:
     for line in fp:
         labels_list.append(line[:-1])
 
-with open(r'ade_palette.txt', 'r') as fp:
+with open(r"ade_palette.txt", "r") as fp:
     for line in fp:
-      tmp_list = list(map(int, line[:-1].strip('][').split(', ')))
-      ade_palette.append(tmp_list)
+        tmp_list = list(map(int, line[:-1].strip("][").split(", ")))
+        ade_palette.append(tmp_list)
 
 colormap = np.asarray(ade_palette)
 
-model_filename = 'segformer-b5-finetuned-ade-640-640.onnx'
-sess = ort.InferenceSession(model_filename)
+model_filename = "segformer-b5-finetuned-ade-640-640.onnx"
+sess_options = ort.SessionOptions()
+sess_options.intra_op_num_threads = os.cpu_count()
+sess = ort.InferenceSession(
+    model_filename, sess_options, providers=["CPUExecutionProvider"]
+)
 
 
 def label_to_color_image(label):
@@ -37,6 +42,7 @@ def label_to_color_image(label):
 
     return colormap[label]
 
+
 def draw_plot(pred_img, seg):
     fig = plt.figure(figsize=(20, 15))
 
@@ -44,7 +50,7 @@ def draw_plot(pred_img, seg):
 
     plt.subplot(grid_spec[0])
     plt.imshow(pred_img)
-    plt.axis('off')
+    plt.axis("off")
 
     LABEL_NAMES = np.asarray(labels_list)
     FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
@@ -59,19 +65,20 @@ def draw_plot(pred_img, seg):
     ax.tick_params(width=0.0, labelsize=25)
     return fig
 
+
 def sepia(input_img):
     img = cv2.imread(input_img)
     img = cv2.resize(img, (640, 640)).astype(np.float32)
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     img_batch = np.expand_dims(img, axis=0)
     img_batch = np.transpose(img_batch, (0, 3, 1, 2))
-    
+
     logits = sess.run(None, {"pixel_values": img_batch})[0]
 
     logits = np.transpose(logits, (0, 2, 3, 1))
-    seg = np.argmax(logits, axis=-1)[0].astype('float32')
-    seg = cv2.resize(seg, (640, 640)).astype('uint8')
-    
+    seg = np.argmax(logits, axis=-1)[0].astype("float32")
+    seg = cv2.resize(seg, (640, 640)).astype("uint8")
+
     color_seg = np.zeros(
         (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
     )  # height, width, 3
@@ -84,11 +91,12 @@ def sepia(input_img):
 
     # Show image + mask
     pred_img = img * 0.5 + color_seg * 0.5
-    pred_img = pred_img.astype(np.uint8)    
+    pred_img = pred_img.astype(np.uint8)
 
     fig = draw_plot(pred_img, seg)
     return fig
 
+
 title = "SegFormer(ADE20k) in TensorFlow"
 description = """
 
@@ -96,12 +104,14 @@ This is demo TensorFlow SegFormer from 🤗 `transformers` official package. The
 
 """
 
-demo = gr.Interface(sepia, 
-                    gr.inputs.Image(type="filepath"), 
-                    outputs=['plot'], 
-                    examples=["ADE_val_00000001.jpeg"],
-                    allow_flagging='never',
-                    title=title,
-                    description=description)
-
-demo.launch()
+demo = gr.Interface(
+    sepia,
+    gr.inputs.Image(type="filepath"),
+    outputs=["plot"],
+    examples=["ADE_val_00000001.jpeg"],
+    allow_flagging="never",
+    title=title,
+    description=description,
+)
+
+demo.launch()