Upload Colony_Analyzer_AI2_HF.py

Colony_Analyzer_AI2_HF.py

@@ -71,10 +71,20 @@ from transformers import SegformerForSemanticSegmentation, SegformerImageProcess
 image_processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b3-finetuned-cityscapes-1024-1024")
 
 def preprocess_image(image):
-    #image = image.convert("RGB")  # Open and convert to RGB
-    inputs = image_processor(image, return_tensors="pt")  # Preprocess for model
+    # Gradio gives image in RGB as numpy array; explicitly ensure it's uint8
+    if image.dtype != np.uint8:
+        image = (255 * image).astype(np.uint8)
+
+    # explicitly confirm 3 channels RGB
+    if image.ndim == 2:
+        image = np.stack([image] * 3, axis=-1)  # grayscale to RGB explicitly
+    elif image.shape[2] == 4:
+        image = image[:, :, :3]  # RGBA to RGB explicitly
+
+    # Now explicitly preprocess using Hugging Face image_processor
+    inputs = image_processor(image, return_tensors="pt")
     return image, inputs["pixel_values"]
-
+    
 def postprocess_mask(logits):
     mask = torch.argmax(logits, dim=1)  # Take argmax across the class dimension
     return mask.squeeze().cpu().numpy()  # Convert to NumPy array