truens66 commited on
Commit
8fcba49
·
verified ·
1 Parent(s): 2d5fbc5

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +4 -2
pipeline.py CHANGED
@@ -2,6 +2,7 @@ import cv2
2
  import numpy as np
3
  import tensorflow as tf
4
  import zipfile
 
5
 
6
  # Set random seed for reproducibility.
7
  tf.random.set_seed(42)
@@ -12,8 +13,9 @@ zip_ref = zipfile.ZipFile(local_zip, 'r')
12
  zip_ref.extractall()
13
  zip_ref.close()
14
 
15
- # Load model
16
- model = tf.keras.models.load_model("efficientnet-b0/")
 
17
 
18
  class DetectionPipeline:
19
  def __init__(self, batch_size=1):
 
2
  import numpy as np
3
  import tensorflow as tf
4
  import zipfile
5
+ import tensorflow_addons as tfa
6
 
7
  # Set random seed for reproducibility.
8
  tf.random.set_seed(42)
 
13
  zip_ref.extractall()
14
  zip_ref.close()
15
 
16
+ # Load model with custom objects
17
+ custom_objects = {"RectifiedAdam": tfa.optimizers.RectifiedAdam}
18
+ model = tf.keras.models.load_model("efficientnet-b0/", custom_objects=custom_objects)
19
 
20
  class DetectionPipeline:
21
  def __init__(self, batch_size=1):