Update pipeline.py

pipeline.py

@@ -46,36 +46,34 @@ class PreTrainedPipeline(Pipeline):
         # Resize image to expected size
 
         expected_input_size = self.model.input_shape
+
+        
+        with Image.open(inputs) as im:
+            inputs = np.array(im)
         if expected_input_size[-1] == 1:  
             inputs = inputs.convert("L")
 
+
+        
         target_size = (expected_input_size[1], expected_input_size[2])
         img = tf.image.resize(inputs, target_size)
-
+        
         img_array = tf.keras.preprocessing.image.img_to_array(img)
         img_array = img_array[tf.newaxis, ...]
-
-        predictions = self.model.predict(img_array, axis=-1)
+        
+        predictions = self.model.predict(img_array)
 
         self.single_output_unit = (
             self.model.output_shape[1] == 1
         )  # if there are two classes
         
-        
-        if self.single_output_unit:
-            score = predictions[0][0]
-            labels = [{"label":"pet", "score":1.0}, {"label":"other", "score":1.0}]
-            #labels = [
-            #    {"label": str(self.id2label["1"]), "score": float(score)},
-            #    {"label": str(self.id2label["0"]), "score": float(1 - score)},
-            ]
-        else:
-            labels = [
-                {
-                    "label": str(self.id2label[str(i)]),
-                    "mask": base64.b64encode(predictions[0][i]),
-                    "score": float(score),
-                }
-                for i, score in enumerate(predictions[0])
-            ]
-        return sorted(labels, key=lambda tup: tup["score"], reverse=True)[: self.top_k]
+        labels = []
+        for i in enumerate(predictions):
+            
+            labels.append({
+                "label": str(i[0]),
+                "mask": base64.b64encode(i[1]),
+                "score": 1.0,
+            })
+
+        return sorted(labels)