Update pipeline.py

pipeline.py

@@ -297,9 +297,7 @@ class PreTrainedPipeline():
             model=self.model2, tokenizer=self.tokenizer2)
     
     def __call__(self, inputs: str)-> List[Dict[str, Any]]:
-        url = 'https://huggingface.co/datasets/mishig/sample_images/resolve/main/cat-1.jpg'
 
-        inputs2 = Image.open(requests.get(url, stream=True).raw)
         # TEMP testing
         # data = [{"video_id": "pqh4LfPeCYs", "start": 835.933, "end": 927.581, "category": "sponsor"}]
         # words = get_words("pqh4LfPeCYs")
@@ -310,70 +308,9 @@ class PreTrainedPipeline():
             'end': 927.581,
             'category': 'sponsor'
         }]
-        segment = data #self.pipeline2(data)
         # END TEMP
 
-
-        # convert img to numpy array, resize and normalize to make the prediction
-        img = np.array(inputs2)
-
-        im = tf.image.resize(img, (128, 128))
-        im = tf.cast(im, tf.float32) / 255.0
-        pred_mask = self.model.predict(im[tf.newaxis, ...])
-        
-        # take the best performing class for each pixel
-        # the output of argmax looks like this [[1, 2, 0], ...]
-        pred_mask_arg = tf.argmax(pred_mask, axis=-1)
-
-        labels = []
-        
-        # convert the prediction mask into binary masks for each class
-        binary_masks = {}
-        mask_codes = {}
-        
-        # when we take tf.argmax() over pred_mask, it becomes a tensor object
-        # the shape becomes TensorShape object, looking like this TensorShape([128]) 
-        # we need to take get shape, convert to list and take the best one
-        
-        rows = pred_mask_arg[0][1].get_shape().as_list()[0]
-        cols = pred_mask_arg[0][2].get_shape().as_list()[0]
-        
-        for cls in range(pred_mask.shape[-1]):
-
-            binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2])) #create masks for each class
-            
-            for row in range(rows):
-
-                for col in range(cols):
-
-                    if pred_mask_arg[0][row][col] == cls:
-                        
-                        binary_masks[f"mask_{cls}"][row][col] = 1
-                    else:
-                        binary_masks[f"mask_{cls}"][row][col] = 0
-
-            mask = binary_masks[f"mask_{cls}"]
-            mask *= 255
-            img = Image.fromarray(mask.astype(np.int8), mode="L")
-               
-            # we need to make it readable for the widget
-            with io.BytesIO() as out:
-                img.save(out, format="PNG")
-                png_string = out.getvalue()
-                mask = base64.b64encode(png_string).decode("utf-8")
-
-            mask_codes[f"mask_{cls}"] = mask
-    
-
-            # widget needs the below format, for each class we return label and mask string
-            labels.append({
-                'inputs': inputs,
-                "label": f"LABEL_{cls}",
-                "mask": mask_codes[f"mask_{cls}"],
-                "score": 1.0,
-                "words": segment
-            })
-        return labels
+        return self.pipeline2(data)