Sophie98 commited on
Commit
12e61cd
1 Parent(s): 7bebb02

change model loading

Browse files
Files changed (1) hide show
  1. StyleTransfer/styleTransfer.py +2 -4
StyleTransfer/styleTransfer.py CHANGED
@@ -77,10 +77,9 @@ def StyleTransformer(content_img: Image.Image, style_img: Image.Image) -> Image.
77
  return Image.fromarray(output)
78
 
79
  ############################################## STYLE-FAST #############################################
 
80
 
81
  def StyleFAST(content_image:Image.Image, style_image:Image.Image) -> Image.Image:
82
- style_transfer_model = tfhub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")
83
-
84
  content_image = tf.convert_to_tensor(np.array(content_image), np.float32)[tf.newaxis, ...] / 255.
85
  style_image = tf.convert_to_tensor(np.array(style_image), np.float32)[tf.newaxis, ...] / 255.
86
  output = style_transfer_model(content_image, style_image)
@@ -88,9 +87,8 @@ def StyleFAST(content_image:Image.Image, style_image:Image.Image) -> Image.Image
88
  return Image.fromarray(np.uint8(stylized_image[0] * 255))
89
 
90
  ########################################### STYLE PROJECTION ##########################################
91
-
92
  def StyleProjection(content_image:Image.Image,style_image:Image.Image) -> Image.Image:
93
- stylepro_artistic = phub.Module(name="stylepro_artistic")
94
  result = stylepro_artistic.style_transfer(
95
  images=[{
96
  'content': np.array(content_image.convert('RGB') )[:, :, ::-1],
77
  return Image.fromarray(output)
78
 
79
  ############################################## STYLE-FAST #############################################
80
+ style_transfer_model = tfhub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")
81
 
82
  def StyleFAST(content_image:Image.Image, style_image:Image.Image) -> Image.Image:
 
 
83
  content_image = tf.convert_to_tensor(np.array(content_image), np.float32)[tf.newaxis, ...] / 255.
84
  style_image = tf.convert_to_tensor(np.array(style_image), np.float32)[tf.newaxis, ...] / 255.
85
  output = style_transfer_model(content_image, style_image)
87
  return Image.fromarray(np.uint8(stylized_image[0] * 255))
88
 
89
  ########################################### STYLE PROJECTION ##########################################
90
+ stylepro_artistic = phub.Module(name="stylepro_artistic")
91
  def StyleProjection(content_image:Image.Image,style_image:Image.Image) -> Image.Image:
 
92
  result = stylepro_artistic.style_transfer(
93
  images=[{
94
  'content': np.array(content_image.convert('RGB') )[:, :, ::-1],