Narsil HF staff commited on
Commit
8b8c024
1 Parent(s): 6d23c1f

model -> self.model.

Browse files
Files changed (1) hide show
  1. pipeline.py +2 -2
pipeline.py CHANGED
@@ -24,7 +24,7 @@ class PreTrainedPipeline():
24
 
25
  im = tf.image.resize(img, (128, 128))
26
  im = tf.cast(im, tf.float32) / 255.0
27
- pred_mask = model.predict(im[tf.newaxis, ...])
28
 
29
  # take the best performing class for each pixel
30
  # the output of argmax looks like this [[1, 2, 0], ...]
@@ -76,4 +76,4 @@ class PreTrainedPipeline():
76
  "mask": mask_codes[f"mask_{cls}"],
77
  "score": 1.0,
78
  })
79
- return labels
24
 
25
  im = tf.image.resize(img, (128, 128))
26
  im = tf.cast(im, tf.float32) / 255.0
27
+ pred_mask = self.model.predict(im[tf.newaxis, ...])
28
 
29
  # take the best performing class for each pixel
30
  # the output of argmax looks like this [[1, 2, 0], ...]
76
  "mask": mask_codes[f"mask_{cls}"],
77
  "score": 1.0,
78
  })
79
+ return labels