Update app.py

app.py

@@ -98,7 +98,7 @@ def foward_pass(image_input: np.ndarray, points: List[List[int]]) -> np.ndarray:
     masks = processor.image_processor.post_process_masks(
         outputs.pred_masks.cpu(), inputs["original_sizes"].to(device), inputs["reshaped_input_sizes"].to(device)
     )
-    masks = masks[0].squeeze(0).numpy().transpose(1, 2, 0)
+    masks = masks[0][0].squeeze(0).numpy()
 
     return masks
 
@@ -114,7 +114,7 @@ def main_func(inputs):
         final_mask = masks[0]
         mask_colors = np.zeros((final_mask.shape[0], final_mask.shape[1], 3), dtype=np.uint8)
         mask_colors[final_mask, :] = np.array([[128, 0, 0]])
-        return Image.fromarray((mask_colors * 0.6 + image_input * 0.4).astype('uint8'), 'RGB')
+        return Image.fromarray((mask_colors+ image_input).astype('uint8'), 'RGB')
     else:
         print('Prediction:: No vehicle found in the image')
         return Image.fromarray(image_input)