nightfury commited on
Commit
5afeae8
·
1 Parent(s): a8d228d

Update colorization.py

Browse files
Files changed (1) hide show
  1. colorization.py +9 -4
colorization.py CHANGED
@@ -8,7 +8,7 @@ from typing import Optional
8
 
9
  model = hub.Module(name='deoldify')
10
  # NOTE: Max is 45 with 11GB video cards. 35 is a good default
11
- render_factor=35
12
 
13
  def colorize_image(image):
14
  # now = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
@@ -23,6 +23,7 @@ def colorize_image(image):
23
  def inferenceColorize(img, version, scale):
24
  # weight /= 100
25
  print(img, scale)
 
26
  try:
27
  extension = os.path.splitext(os.path.basename(str(img)))[1]
28
  img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
@@ -55,17 +56,21 @@ def inferenceColorize(img, version, scale):
55
  extension = 'png'
56
  else:
57
  extension = 'jpg'
58
-
 
59
  model.predict(img.name)
60
  save_path = f'output/DeOldify/'+Path(img.name).stem+'.{extension}'
61
 
 
62
  cv2.imwrite(save_path, output)
63
  output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
64
- return output, save_path
 
 
65
  except Exception as error:
66
  print('global exception', error)
67
  return None, None
68
-
69
 
70
  colorize = gr.Interface(
71
  inferenceColorize, [
 
8
 
9
  model = hub.Module(name='deoldify')
10
  # NOTE: Max is 45 with 11GB video cards. 35 is a good default
11
+ render_factor=40
12
 
13
  def colorize_image(image):
14
  # now = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
 
23
  def inferenceColorize(img, version, scale):
24
  # weight /= 100
25
  print(img, scale)
26
+ '''
27
  try:
28
  extension = os.path.splitext(os.path.basename(str(img)))[1]
29
  img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
 
56
  extension = 'png'
57
  else:
58
  extension = 'jpg'
59
+ '''
60
+ os.makedirs('output', exist_ok=True)
61
  model.predict(img.name)
62
  save_path = f'output/DeOldify/'+Path(img.name).stem+'.{extension}'
63
 
64
+ '''
65
  cv2.imwrite(save_path, output)
66
  output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
67
+ '''
68
+ return img, save_path
69
+ '''
70
  except Exception as error:
71
  print('global exception', error)
72
  return None, None
73
+ '''
74
 
75
  colorize = gr.Interface(
76
  inferenceColorize, [