Upload folder using huggingface_hub

README.md

@@ -54,13 +54,18 @@ result = segmenter("ultrasound.png", organ_id=3)
 If `organ_id` is not provided, the model automatically uses `-1`, matching the
 unknown-organ conditioning used in training.
 
+Known organ IDs are: appendix `0`, breast `1`, cardiac `2`, thyroid `3`, fetal
+`4`, kidney `5`, liver `6`, and testicle `7`. To reproduce the original
+evaluation pipeline for a Fetal HC image, use `organ_id=4`.
+
 ## Output
 
 The pipeline applies a sigmoid to the foreground logit and returns a binary
-PIL mask thresholded at 0.5. Change the threshold if needed:
+PIL mask thresholded at 0.7, matching the original evaluation pipeline. Change
+the threshold if needed:
 
 ```python
-result = segmenter("ultrasound.png", threshold=0.4)
+result = segmenter("ultrasound.png", threshold=0.7)
 ```
 
 ## Preprocessing

image_processing_cond_unet.py

@@ -2,8 +2,8 @@ from typing import Optional, Union
 
 import numpy as np
 import torch
-import torch.nn.functional as F
 from PIL import Image
+from torchvision.transforms import v2
 from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
 
 
@@ -51,32 +51,22 @@ class CondUNetImageProcessor(BaseImageProcessor):
             image = image.expand(3, -1, -1)
         if image.shape[0] != 3:
             raise ValueError("Cond-UNet requires one or three input channels.")
-        image = image.to(dtype=torch.float32)
-        if image.max() <= 1:
-            image = image * 255.0
         height, width = image.shape[-2:]
         if self.keep_aspect_ratio:
-            scale = self.image_size / max(height, width)
-            new_height = int(height * scale)
-            new_width = int(width * scale)
+            resize_factor = max(height, width) / self.image_size
+            new_height = int(height / resize_factor)
+            new_width = int(width / resize_factor)
             new_height += new_height % 2
             new_width += new_width % 2
-            image = F.interpolate(
-                image.unsqueeze(0),
-                size=(new_height, new_width),
-                mode="bilinear",
-                align_corners=False,
-            ).squeeze(0)
+            image = v2.functional.resize(image, [new_height, new_width])
             pad_left = (self.image_size - new_width) // 2
             pad_top = (self.image_size - new_height) // 2
-            image = F.pad(image, (pad_left, self.image_size - new_width - pad_left, pad_top, self.image_size - new_height - pad_top))
+            image = v2.functional.pad(image, fill=0, padding=[pad_left, pad_top])
         else:
-            image = F.interpolate(
-                image.unsqueeze(0),
-                size=(self.image_size, self.image_size),
-                mode="bilinear",
-                align_corners=False,
-            ).squeeze(0)
+            image = v2.functional.resize(image, [self.image_size, self.image_size])
+        image = image.to(dtype=torch.float32)
+        if image.max() <= 1:
+            image = image * 255.0
         mean = torch.tensor(self.mean, dtype=image.dtype).view(-1, 1, 1)
         std = torch.tensor(self.std, dtype=image.dtype).view(-1, 1, 1)
         return (image - mean) / std

pipeline.py

@@ -42,9 +42,11 @@ class CondUNetImageSegmentationPipeline(Pipeline):
         outputs = self.model(**model_inputs)
         return {"logits": outputs.logits, "original_size": original_size}
 
-    def postprocess(self, model_outputs, threshold=0.5, **kwargs):
+    def postprocess(self, model_outputs, threshold=0.7, **kwargs):
         logits = model_outputs["logits"]
         height, width = model_outputs["original_size"]
-        probabilities = torch.sigmoid(F.interpolate(logits, size=(height, width), mode="bilinear", align_corners=False))[0, 0]
+        probabilities = torch.sigmoid(
+            F.interpolate(logits, size=(height, width), mode="nearest")
+        )[0, 0]
         mask = (probabilities >= threshold).to(torch.uint8).cpu().numpy() * 255
         return {"label": "foreground", "mask": Image.fromarray(mask), "score": float(probabilities.mean())}