artelabsuper commited on
Commit
7f268fe
1 Parent(s): 63cfb07

info on demo and test script use gpu

Browse files
Files changed (2) hide show
  1. app.py +4 -2
  2. test.py +2 -2
app.py CHANGED
@@ -51,8 +51,9 @@ def predict(input_image, model_name):
51
  fig.canvas.draw()
52
  data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
53
  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
54
- # return correct image
55
- return sr, data
 
56
 
57
  iface = gr.Interface(
58
  fn=predict,
@@ -61,6 +62,7 @@ iface = gr.Interface(
61
  gr.inputs.Radio(MODELS_TYPE)
62
  ],
63
  outputs=[
 
64
  gr.Image(),
65
  gr.Image()
66
  ],
 
51
  fig.canvas.draw()
52
  data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
53
  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
54
+ # return correct image and info
55
+ info = f"{model_name} with {sum(p.numel() for p in generators[MODELS_TYPE.index(model_name)].parameters())} parameters"
56
+ return info, sr, data
57
 
58
  iface = gr.Interface(
59
  fn=predict,
 
62
  gr.inputs.Radio(MODELS_TYPE)
63
  ],
64
  outputs=[
65
+ gr.Text(),
66
  gr.Image(),
67
  gr.Image()
68
  ],
test.py CHANGED
@@ -10,8 +10,8 @@ from models.modelNetC import Generator as GC
10
 
11
 
12
 
13
- DEVICE='cpu'
14
- model_type = 'model_b'
15
 
16
  modeltype2path = {
17
  'model_a': 'DTM_exp_train10%_model_a/g-best.pth',
 
10
 
11
 
12
 
13
+ DEVICE='cuda'
14
+ model_type = 'model_c'
15
 
16
  modeltype2path = {
17
  'model_a': 'DTM_exp_train10%_model_a/g-best.pth',