rwheel commited on
Commit
314893c
1 Parent(s): 4ee69d5

Use kornia io module to load the image

Browse files
Files changed (1) hide show
  1. app.py +10 -16
app.py CHANGED
@@ -1,61 +1,55 @@
1
  import gradio as gr
2
 
3
- import cv2
4
- import numpy as np
5
-
6
- import torch
7
  import kornia as K
8
  from kornia.core import Tensor
9
 
10
  def edge_detection(filepath, detector):
11
 
12
- img_bgr: np.ndarray = cv2.imread(filepath, cv2.IMREAD_COLOR)
13
-
14
- x_bgr: torch.Tensor = K.utils.image_to_tensor(img_bgr)
15
- x_bgr = x_bgr[None,...].float() / 255.
16
 
17
- x_rgb: torch.Tensor = K.color.bgr_to_rgb(x_bgr)
18
  x_gray = K.color.rgb_to_grayscale(x_rgb)
19
 
20
 
21
  if detector == '1st order derivates in x':
22
- grads: torch.Tensor = K.filters.spatial_gradient(x_gray, order=1)
23
  grads_x = grads[:, :, 0]
24
  grads_y = grads[:, :, 1]
25
 
26
  output = K.utils.tensor_to_image(1. - grads_x.clamp(0., 1.))
27
 
28
  elif detector == '1st order derivates in y':
29
- grads: torch.Tensor = K.filters.spatial_gradient(x_gray, order=1)
30
  grads_x = grads[:, :, 0]
31
  grads_y = grads[:, :, 1]
32
 
33
  output = K.utils.tensor_to_image(1. - grads_y.clamp(0., 1.))
34
 
35
  elif detector == '2nd order derivatives in x':
36
- grads: torch.Tensor = K.filters.spatial_gradient(x_gray, order=2)
37
  grads_x = grads[:, :, 0]
38
  grads_y = grads[:, :, 1]
39
 
40
  output = K.utils.tensor_to_image(1. - grads_x.clamp(0., 1.))
41
 
42
  elif detector == '2nd order derivatives in y':
43
- grads: torch.Tensor = K.filters.spatial_gradient(x_gray, order=2)
44
  grads_x = grads[:, :, 0]
45
  grads_y = grads[:, :, 1]
46
 
47
  output = K.utils.tensor_to_image(1. - grads_y.clamp(0., 1.))
48
 
49
  elif detector == 'Sobel':
50
- x_sobel: torch.Tensor = K.filters.sobel(x_gray)
51
  output = K.utils.tensor_to_image(1. - x_sobel)
52
 
53
  elif detector == 'Laplacian':
54
- x_laplacian: torch.Tensor = K.filters.laplacian(x_gray, kernel_size=5)
55
  output = K.utils.tensor_to_image(1. - x_laplacian.clamp(0., 1.))
56
 
57
  else:
58
- x_canny: torch.Tensor = K.filters.canny(x_gray)[0]
59
  output = K.utils.tensor_to_image(1. - x_canny.clamp(0., 1.0))
60
 
61
  return output
 
1
  import gradio as gr
2
 
 
 
 
 
3
  import kornia as K
4
  from kornia.core import Tensor
5
 
6
  def edge_detection(filepath, detector):
7
 
8
+ img: Tensor = K.io.load_image(filepath, K.io.ImageLoadType.RGB32)
9
+ img = img[None]
 
 
10
 
11
+ x_rgb: Tensor = K.color.bgr_to_rgb(img)
12
  x_gray = K.color.rgb_to_grayscale(x_rgb)
13
 
14
 
15
  if detector == '1st order derivates in x':
16
+ grads: Tensor = K.filters.spatial_gradient(x_gray, order=1)
17
  grads_x = grads[:, :, 0]
18
  grads_y = grads[:, :, 1]
19
 
20
  output = K.utils.tensor_to_image(1. - grads_x.clamp(0., 1.))
21
 
22
  elif detector == '1st order derivates in y':
23
+ grads: Tensor = K.filters.spatial_gradient(x_gray, order=1)
24
  grads_x = grads[:, :, 0]
25
  grads_y = grads[:, :, 1]
26
 
27
  output = K.utils.tensor_to_image(1. - grads_y.clamp(0., 1.))
28
 
29
  elif detector == '2nd order derivatives in x':
30
+ grads: Tensor = K.filters.spatial_gradient(x_gray, order=2)
31
  grads_x = grads[:, :, 0]
32
  grads_y = grads[:, :, 1]
33
 
34
  output = K.utils.tensor_to_image(1. - grads_x.clamp(0., 1.))
35
 
36
  elif detector == '2nd order derivatives in y':
37
+ grads: Tensor = K.filters.spatial_gradient(x_gray, order=2)
38
  grads_x = grads[:, :, 0]
39
  grads_y = grads[:, :, 1]
40
 
41
  output = K.utils.tensor_to_image(1. - grads_y.clamp(0., 1.))
42
 
43
  elif detector == 'Sobel':
44
+ x_sobel: Tensor = K.filters.sobel(x_gray)
45
  output = K.utils.tensor_to_image(1. - x_sobel)
46
 
47
  elif detector == 'Laplacian':
48
+ x_laplacian: Tensor = K.filters.laplacian(x_gray, kernel_size=5)
49
  output = K.utils.tensor_to_image(1. - x_laplacian.clamp(0., 1.))
50
 
51
  else:
52
+ x_canny: Tensor = K.filters.canny(x_gray)[0]
53
  output = K.utils.tensor_to_image(1. - x_canny.clamp(0., 1.0))
54
 
55
  return output