mischeiwiller commited on
Commit
bdf18f8
1 Parent(s): 68db27f

fix: handle both numpy arrays and file paths in inference function

Browse files
Files changed (1) hide show
  1. app.py +13 -7
app.py CHANGED
@@ -3,12 +3,18 @@ import cv2
3
  import matplotlib.pyplot as plt
4
  import numpy as np
5
  import torch
6
- import torchvision
7
  import kornia as K
8
 
9
  def inference(file1,num_iters):
10
- img: np.ndarray = cv2.imread(file1.name)
11
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
 
 
 
 
 
 
 
12
  img = img + np.random.normal(loc=0.0, scale=0.1, size=img.shape)
13
  img = np.clip(img, 0.0, 1.0)
14
 
@@ -55,13 +61,13 @@ examples = [ ["doraemon.png",2000]
55
 
56
 
57
  inputs = [
58
- gr.Image(type='file', label='Input Image'),
59
- gr.Slider(minimum=50, maximum=10000, step=50, default=500, label="num_iters")
60
  ]
61
 
62
  outputs = [
63
- gr.Image(type='file', label='Noised Image'),
64
- gr.Image(type='file', label='Denoised Image'),
65
  ]
66
 
67
  title = "Denoise image using total variation"
 
3
  import matplotlib.pyplot as plt
4
  import numpy as np
5
  import torch
 
6
  import kornia as K
7
 
8
  def inference(file1,num_iters):
9
+ # Check if file1 is already a numpy array
10
+ if isinstance(file1, np.ndarray):
11
+ img = file1
12
+ else:
13
+ # If it's not a numpy array, assume it's a file path
14
+ img = cv2.imread(file1)
15
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
16
+
17
+ img = img.astype(np.float32) / 255.0
18
  img = img + np.random.normal(loc=0.0, scale=0.1, size=img.shape)
19
  img = np.clip(img, 0.0, 1.0)
20
 
 
61
 
62
 
63
  inputs = [
64
+ gr.Image(type='numpy', label='Input Image'),
65
+ gr.Slider(minimum=50, maximum=10000, step=50, value=500, label="num_iters")
66
  ]
67
 
68
  outputs = [
69
+ gr.Image(type='numpy', label='Noised Image'),
70
+ gr.Image(type='numpy', label='Denoised Image'),
71
  ]
72
 
73
  title = "Denoise image using total variation"