Spaces:
Running
Running
| """ | |
| Module to manage the segmentation YOLO model. | |
| """ | |
| from pathlib import Path | |
| from typing import Any | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| from ultralytics import YOLO | |
| import yolo | |
| def load_pretrained_model(model_str: str) -> YOLO: | |
| """ | |
| Load the pretrained model. | |
| """ | |
| return yolo.load_pretrained_model(model_str) | |
| def train( | |
| model: YOLO, | |
| data_yaml_path: Path, | |
| params: dict, | |
| project: Path = Path("data/04_models/yolo/"), | |
| experiment_name: str = "train", | |
| ): | |
| """Main function for running a train run. It saves the results | |
| under `project / experiment_name`. | |
| Args: | |
| model (YOLO): result of `load_pretrained_model`. | |
| data_yaml_path (Path): filepath to the data.yaml file that specifies the split and classes to train on | |
| params (dict): parameters to override when running the training. See https://docs.ultralytics.com/modes/train/#train-settings for a complete list of parameters. | |
| project (Path): root path to store the run artifacts and results. | |
| experiment_name (str): name of the experiment, that is added to the project root path to store the run. | |
| """ | |
| return yolo.train( | |
| model=model, | |
| data_yaml_path=data_yaml_path, | |
| params=params, | |
| project=project, | |
| experiment_name=experiment_name, | |
| ) | |
| def predict(model: YOLO, pil_image: Image.Image) -> dict[str, Any]: | |
| """ | |
| Given a loaded model an a PIL image, it returns a map | |
| containing the segmentation predictions. | |
| Args: | |
| model (YOLO): loaded YOLO model for segmentation. | |
| pil_image (PIL): image to run the model on. | |
| Returns: | |
| prediction: Raw prediction from the model. | |
| orig_image: original image used for inference | |
| after preprocessing stages applied by | |
| ultralytics. | |
| mask (PIL): postprocessed mask in white and black format - used for visualization | |
| mask_raw (np.ndarray): Raw mask not postprocessed | |
| masked (PIL): mask applied to the pil_image. | |
| """ | |
| predictions = model(pil_image) | |
| mask_raw = predictions[0].masks[0].data.cpu().numpy().transpose(1, 2, 0).squeeze() | |
| # Convert single channel grayscale to 3 channel image | |
| mask_3channel = cv2.merge((mask_raw, mask_raw, mask_raw)) | |
| # Get the size of the original image (height, width, channels) | |
| h2, w2, c2 = predictions[0].orig_img.shape | |
| # Resize the mask to the same size as the image (can probably be removed if image is the same size as the model) | |
| mask = cv2.resize(mask_3channel, (w2, h2)) | |
| # Convert BGR to HSV | |
| hsv = cv2.cvtColor(mask, cv2.COLOR_BGR2HSV) | |
| # Define range of brightness in HSV | |
| lower_black = np.array([0, 0, 0]) | |
| upper_black = np.array([0, 0, 1]) | |
| # Create a mask. Threshold the HSV image to get everything black | |
| mask = cv2.inRange(mask, lower_black, upper_black) | |
| # Invert the mask to get everything but black | |
| mask = cv2.bitwise_not(mask) | |
| # Apply the mask to the original image | |
| masked = cv2.bitwise_and( | |
| predictions[0].orig_img, | |
| predictions[0].orig_img, | |
| mask=mask, | |
| ) | |
| # bgr to rgb and PIL conversion | |
| image_output2 = Image.fromarray(masked[:, :, ::-1]) | |
| # return Image.fromarray(mask), image_output2 | |
| return { | |
| "prediction": predictions[0], | |
| "mask": Image.fromarray(mask), | |
| "mask_raw": mask_raw, | |
| "masked": Image.fromarray(masked[:, :, ::-1]), | |
| } | |