Merge pull request #13 from Imageomics/12-rgba-fix

Fix for RGBA png files

Author: johnbradley

src/bioclip/predict.py

@@ -119,6 +119,11 @@ def get_txt_names():
     return txt_names
 
 
+def open_image(image_path):
+    img = PIL.Image.open(image_path)
+    return img.convert("RGB")
+
+
 preprocess_img = transforms.Compose(
     [
         transforms.ToTensor(),
@@ -181,7 +186,7 @@ def get_txt_features(self, classnames):
 
     @torch.no_grad()
     def predict(self, image_path: str, cls_ary: List[str]) -> dict[str, float]:
-        img = PIL.Image.open(image_path)
+        img = open_image(image_path)
         classes = [cls.strip() for cls in cls_ary]
         txt_features = self.get_txt_features(classes)
 
@@ -248,7 +253,7 @@ def __init__(self, device: Union[str, torch.device] = 'cpu', model_str: str = MO
 
     @torch.no_grad()
     def get_image_features(self, image_path: str) -> torch.Tensor:
-        img = PIL.Image.open(image_path)
+        img = open_image(image_path)
         return self.encode_image(img)
 
     def encode_image(self, img: PIL.Image.Image) -> torch.Tensor:
@@ -295,7 +300,7 @@ def format_grouped_probs(self, image_path: str, probs: torch.Tensor, rank: Rank,
 
     @torch.no_grad()
     def predict(self, image_path: str, rank: Rank, min_prob: float = 1e-9, k: int = 5) -> List[dict[str, float]]:
-        img = PIL.Image.open(image_path)
+        img = open_image(image_path)
         probs = self.predict_species(img)
         if rank == Rank.SPECIES:
             return self.format_species_probs(image_path, probs, k)

tests/images/mycat.png

@@ -7,6 +7,7 @@
 
 DIRNAME = os.path.dirname(os.path.realpath(__file__))
 EXAMPLE_CAT_IMAGE = os.path.join(DIRNAME, "images", "mycat.jpg")
+EXAMPLE_CAT_IMAGE2 = os.path.join(DIRNAME, "images", "mycat.png")
 
 class TestPredict(unittest.TestCase):
     def test_tree_of_life_classifier_species(self):
@@ -51,6 +52,12 @@ def test_custom_labels_classifier(self):
             {'file_name': EXAMPLE_CAT_IMAGE, 'classification': 'dog', 'score': unittest.mock.ANY},
         ])
 
+    def test_predict_with_rgba_image(self):
+        # Ensure that the classifier can handle RGBA images
+        classifier = TreeOfLifeClassifier()
+        prediction_ary = classifier.predict(image_path=EXAMPLE_CAT_IMAGE2, rank=Rank.SPECIES)
+        self.assertEqual(len(prediction_ary), 5)
+
 
 class TestEmbed(unittest.TestCase):
     def test_get_image_features(self):