maloyan commited on
Commit
bfb7b59
·
1 Parent(s): 6282546

no dependencies

Browse files
Files changed (1) hide show
  1. app.py +2 -33
app.py CHANGED
@@ -1,8 +1,8 @@
1
  from argparse import Namespace
2
 
3
- import cv2
4
  import gradio as gr
5
- import numpy as np
6
  import torch
7
  import torchvision.transforms as transforms
8
  from huggingface_hub import hf_hub_download
@@ -17,15 +17,6 @@ transfroms = transforms.Compose([
17
  transforms.ToTensor()]
18
  )
19
 
20
- def log_input_image(x, opts):
21
- if opts.label_nc == 0:
22
- return tensor2im(x)
23
- elif opts.label_nc == 1:
24
- return tensor2sketch(x)
25
- else:
26
- return tensor2map(x)
27
-
28
-
29
  def tensor2im(var):
30
  var = var.cpu().detach().transpose(0, 2).transpose(0, 1).numpy()
31
  var = ((var + 1) / 2)
@@ -34,28 +25,6 @@ def tensor2im(var):
34
  var = var * 255
35
  return Image.fromarray(var.astype('uint8'))
36
 
37
- def tensor2map(var):
38
- mask = np.argmax(var.data.cpu().numpy(), axis=0)
39
- print(np.unique(mask))
40
- colors = get_colors()
41
- mask_image = np.zeros(shape=(mask.shape[0], mask.shape[1], 3))
42
- for class_idx in np.unique(mask):
43
- mask_image[mask == class_idx] = colors[class_idx]
44
- mask_image = mask_image.astype('uint8')
45
- return Image.fromarray(mask_image)
46
-
47
- def tensor2sketch(var):
48
- im = var[0].cpu().detach().numpy()
49
- im = cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)
50
- im = (im * 255).astype(np.uint8)
51
- return Image.fromarray(im)
52
- def get_colors():
53
- # currently support up to 19 classes (for the celebs-hq-mask dataset)
54
- colors = [[0, 0, 0], [204, 0, 0], [76, 153, 0], [204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255],
55
- [255, 204, 204], [102, 51, 0], [255, 0, 0], [102, 204, 0], [255, 255, 0], [0, 0, 153], [0, 0, 204],
56
- [255, 51, 153], [0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]]
57
- return colors
58
-
59
  def sketch_recognition(img):
60
  from_im = transfroms(Image.fromarray(img))
61
  with torch.no_grad():
 
1
  from argparse import Namespace
2
 
3
+ #import cv2
4
  import gradio as gr
5
+ #import numpy as np
6
  import torch
7
  import torchvision.transforms as transforms
8
  from huggingface_hub import hf_hub_download
 
17
  transforms.ToTensor()]
18
  )
19
 
 
 
 
 
 
 
 
 
 
20
  def tensor2im(var):
21
  var = var.cpu().detach().transpose(0, 2).transpose(0, 1).numpy()
22
  var = ((var + 1) / 2)
 
25
  var = var * 255
26
  return Image.fromarray(var.astype('uint8'))
27
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
28
  def sketch_recognition(img):
29
  from_im = transfroms(Image.fromarray(img))
30
  with torch.no_grad():