set intra op threads

#5
by sayakpaul HF staff - opened
Files changed (1) hide show
  1. app.py +36 -26
app.py CHANGED
@@ -1,31 +1,36 @@
1
- import gradio as gr
2
-
3
- import sys
4
  import csv
5
- import numpy as np
 
 
6
  import cv2
7
- from matplotlib import gridspec
8
  import matplotlib.pyplot as plt
 
9
  import onnxruntime as ort
 
10
 
11
  ade_palette = []
12
  labels_list = []
13
 
14
  csv.field_size_limit(sys.maxsize)
15
 
16
- with open(r'labels.txt', 'r') as fp:
17
  for line in fp:
18
  labels_list.append(line[:-1])
19
 
20
- with open(r'ade_palette.txt', 'r') as fp:
21
  for line in fp:
22
- tmp_list = list(map(int, line[:-1].strip('][').split(', ')))
23
- ade_palette.append(tmp_list)
24
 
25
  colormap = np.asarray(ade_palette)
26
 
27
- model_filename = 'segformer-b5-finetuned-ade-640-640.onnx'
28
- sess = ort.InferenceSession(model_filename)
 
 
 
 
29
 
30
 
31
  def label_to_color_image(label):
@@ -37,6 +42,7 @@ def label_to_color_image(label):
37
 
38
  return colormap[label]
39
 
 
40
  def draw_plot(pred_img, seg):
41
  fig = plt.figure(figsize=(20, 15))
42
 
@@ -44,7 +50,7 @@ def draw_plot(pred_img, seg):
44
 
45
  plt.subplot(grid_spec[0])
46
  plt.imshow(pred_img)
47
- plt.axis('off')
48
 
49
  LABEL_NAMES = np.asarray(labels_list)
50
  FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
@@ -59,19 +65,20 @@ def draw_plot(pred_img, seg):
59
  ax.tick_params(width=0.0, labelsize=25)
60
  return fig
61
 
 
62
  def sepia(input_img):
63
  img = cv2.imread(input_img)
64
  img = cv2.resize(img, (640, 640)).astype(np.float32)
65
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
66
  img_batch = np.expand_dims(img, axis=0)
67
  img_batch = np.transpose(img_batch, (0, 3, 1, 2))
68
-
69
  logits = sess.run(None, {"pixel_values": img_batch})[0]
70
 
71
  logits = np.transpose(logits, (0, 2, 3, 1))
72
- seg = np.argmax(logits, axis=-1)[0].astype('float32')
73
- seg = cv2.resize(seg, (640, 640)).astype('uint8')
74
-
75
  color_seg = np.zeros(
76
  (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
77
  ) # height, width, 3
@@ -84,11 +91,12 @@ def sepia(input_img):
84
 
85
  # Show image + mask
86
  pred_img = img * 0.5 + color_seg * 0.5
87
- pred_img = pred_img.astype(np.uint8)
88
 
89
  fig = draw_plot(pred_img, seg)
90
  return fig
91
 
 
92
  title = "SegFormer(ADE20k) in TensorFlow"
93
  description = """
94
 
@@ -96,12 +104,14 @@ This is demo TensorFlow SegFormer from 🤗 `transformers` official package. The
96
 
97
  """
98
 
99
- demo = gr.Interface(sepia,
100
- gr.inputs.Image(type="filepath"),
101
- outputs=['plot'],
102
- examples=["ADE_val_00000001.jpeg"],
103
- allow_flagging='never',
104
- title=title,
105
- description=description)
106
-
107
- demo.launch()
 
 
 
 
 
 
1
  import csv
2
+ import os
3
+ import sys
4
+
5
  import cv2
6
+ import gradio as gr
7
  import matplotlib.pyplot as plt
8
+ import numpy as np
9
  import onnxruntime as ort
10
+ from matplotlib import gridspec
11
 
12
  ade_palette = []
13
  labels_list = []
14
 
15
  csv.field_size_limit(sys.maxsize)
16
 
17
+ with open(r"labels.txt", "r") as fp:
18
  for line in fp:
19
  labels_list.append(line[:-1])
20
 
21
+ with open(r"ade_palette.txt", "r") as fp:
22
  for line in fp:
23
+ tmp_list = list(map(int, line[:-1].strip("][").split(", ")))
24
+ ade_palette.append(tmp_list)
25
 
26
  colormap = np.asarray(ade_palette)
27
 
28
+ model_filename = "segformer-b5-finetuned-ade-640-640.onnx"
29
+ sess_options = ort.SessionOptions()
30
+ sess_options.intra_op_num_threads = os.cpu_count()
31
+ sess = ort.InferenceSession(
32
+ model_filename, sess_options, providers=["CPUExecutionProvider"]
33
+ )
34
 
35
 
36
  def label_to_color_image(label):
 
42
 
43
  return colormap[label]
44
 
45
+
46
  def draw_plot(pred_img, seg):
47
  fig = plt.figure(figsize=(20, 15))
48
 
 
50
 
51
  plt.subplot(grid_spec[0])
52
  plt.imshow(pred_img)
53
+ plt.axis("off")
54
 
55
  LABEL_NAMES = np.asarray(labels_list)
56
  FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
 
65
  ax.tick_params(width=0.0, labelsize=25)
66
  return fig
67
 
68
+
69
  def sepia(input_img):
70
  img = cv2.imread(input_img)
71
  img = cv2.resize(img, (640, 640)).astype(np.float32)
72
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
73
  img_batch = np.expand_dims(img, axis=0)
74
  img_batch = np.transpose(img_batch, (0, 3, 1, 2))
75
+
76
  logits = sess.run(None, {"pixel_values": img_batch})[0]
77
 
78
  logits = np.transpose(logits, (0, 2, 3, 1))
79
+ seg = np.argmax(logits, axis=-1)[0].astype("float32")
80
+ seg = cv2.resize(seg, (640, 640)).astype("uint8")
81
+
82
  color_seg = np.zeros(
83
  (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
84
  ) # height, width, 3
 
91
 
92
  # Show image + mask
93
  pred_img = img * 0.5 + color_seg * 0.5
94
+ pred_img = pred_img.astype(np.uint8)
95
 
96
  fig = draw_plot(pred_img, seg)
97
  return fig
98
 
99
+
100
  title = "SegFormer(ADE20k) in TensorFlow"
101
  description = """
102
 
 
104
 
105
  """
106
 
107
+ demo = gr.Interface(
108
+ sepia,
109
+ gr.inputs.Image(type="filepath"),
110
+ outputs=["plot"],
111
+ examples=["ADE_val_00000001.jpeg"],
112
+ allow_flagging="never",
113
+ title=title,
114
+ description=description,
115
+ )
116
+
117
+ demo.launch()