diff --git a/evaluation/python/generate_masks.py b/evaluation/python/generate_masks.py index 217c6b8..927a4b3 100755 --- a/evaluation/python/generate_masks.py +++ b/evaluation/python/generate_masks.py @@ -2,7 +2,7 @@ import numpy as np import tensorflow as tf -from PIL import Image, ImageOps +import cv2 from pathlib import Path SEG_MODEL_FILE_PATH = "../../app/build/downloads/fairscan-segmentation-model.tflite" @@ -12,15 +12,21 @@ INPUT_WIDTH = 256 INPUT_HEIGHT = 256 def get_resized_image(image_path): - img = Image.open(image_path).convert("RGB") - img = ImageOps.exif_transpose(img) - img = img.convert("RGB").resize((INPUT_WIDTH, INPUT_HEIGHT), Image.BILINEAR) + img = cv2.imread(str(image_path), cv2.IMREAD_COLOR) + if img is None: + raise ValueError(f"Failed to load image: {image_path}") + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = cv2.resize( + img, + (INPUT_WIDTH, INPUT_HEIGHT), + interpolation=cv2.INTER_LINEAR + ) return img -def preprocess_image(img: Image.Image) -> np.ndarray: - img_np = np.asarray(img).astype(np.float32) - img_np = (img_np - 127.5) / 127.5 # Normalize to [-1, 1] - return np.expand_dims(img_np, axis=0) +def preprocess_image(img): + img = img.astype(np.float32) + img = img / 127.5 - 1.0 + return img[np.newaxis, ...] def postprocess_output(output: np.ndarray) -> np.ndarray: output = np.squeeze(output).astype(np.float32) # Shape: (256, 256) @@ -45,7 +51,7 @@ mask_input_dir = DATASET_DIR / "masks" for image_path in sorted(img_input_dir.glob("*.jpg")): print(f"Generating mask for {image_path}") img = get_resized_image(image_path) - img = ImageOps.exif_transpose(img) mask = get_segmentation_mask(img) mask_path = mask_input_dir / (image_path.stem + ".png") - Image.fromarray((mask * 255).astype(np.uint8)).save(mask_path) + mask_uint8 = (mask * 255.0).round().astype(np.uint8) + cv2.imwrite(str(mask_path), mask_uint8) diff --git a/evaluation/python/requirements.txt b/evaluation/python/requirements.txt index 4e3df80..fe5cfc8 100644 --- a/evaluation/python/requirements.txt +++ b/evaluation/python/requirements.txt @@ -1,3 +1,3 @@ tensorflow numpy -Pillow +opencv-python