Source code for fortrace.utility.image_processing.opencv_utils

import os
import pathlib
from typing import Literal

import cv2
import numpy as np


[docs] class OpenCVRectangle: def __init__( self, x0: int, y0: int, x1: int, y1: int, coordination_format: Literal["xyxy", "xywh"] = "xyxy", ): """Representation of a OpenCV rectangle. Args: x0: x coordinate of the top left corner of the rectangle y0: y coordinate of the top left corner of the rectangle x1: x coordinate of the bottom right corner of the rectangle y1: y coordinate of the bottom right corner of the rectangle coordination_format: format of the rectangle coordinates Note: p0(x0, y0)----w-------------+ | | | | | h | | | | +-------------------p1(x1,y1) """ if coordination_format == "xyxy": self.x0 = x0 self.y0 = y0 self.x1 = x1 self.y1 = y1 elif coordination_format == "xywh": self.x0 = x0 self.y0 = y0 self.x1 = x0 + x1 self.y1 = y0 + y1 else: raise ValueError(f"Invalid format {coordination_format}")
[docs] def __getitem__(self, item: int | slice) -> int | list[int]: """Enable to use class like tuple or list and access elements with [<item>]. Args: item: index of the element to access (0 <= item <= 4) or slice Returns: Coordinate(s) at given index. """ if isinstance(item, slice): return [ self.__getitem__(i) for i in range( item.start if item.start else 0, item.stop if item.stop else 4, item.step if item.step else 1, ) ] else: if item == 0: return self.x0 elif item == 1: return self.y0 elif item == 2: return self.x1 elif item == 3: return self.y1 else: raise IndexError("Rectangle index (0 <= item <= 4) out of range")
[docs] def p0(self) -> tuple[int, int]: """Upper left point of the rectangle. Note: (0,0)-----------------------+ | image | | p0-------+ | | | rect | | | +--------+ | | | +---------------------------+ Returns: The point p0: (x0, y0) of the rectangle """ return self.x0, self.y0
[docs] def p1(self) -> tuple[int, int]: """Lower right point of the rectangle. Note: (0,0)-----------------------+ | image | | +--------+ | | | rect | | | +-------p1 | | | +---------------------------+ Returns: The point p1 (x1, y1) of the rectangle """ return self.x1, self.y1
@property def w(self): """Width of the rectangle.""" return self.x1 - self.x0 @property def h(self): """Height of the rectangle.""" return self.y1 - self.y0
[docs] def centroid(self) -> tuple[int, int]: """Centroid of the rectangle.""" return (self.x0 + self.x1) // 2, (self.y0 + self.y1) // 2
[docs] def draw_rectangle( image: cv2.Mat | np.ndarray, rect: OpenCVRectangle, inplace: bool = True ) -> cv2.Mat | np.ndarray: """Draw a rectangle on the image. Args: image: captured image, readable by openCV rect: OpenCVRectangle to draw on the image inplace: draw the rectangle on the original image, altering it permanently Returns: openCV image with a drawn on rectangle """ return cv2.rectangle( image if inplace else image.copy(), rect.p0(), rect.p1(), (0, 255, 0), 2 )
[docs] def extract_rectangle( image: cv2.Mat | np.ndarray, rect: OpenCVRectangle ) -> cv2.Mat | np.ndarray: """Extract the area of the rectangle from the image. Args: image: OpenCV image rect: rectangle to be extracted Returns: extracted sub-area of the image """ return image[rect.y1 : rect.y0, rect.x0 : rect.x1]
[docs] def show_image( image: np.ndarray, title: str = "title", rectangle: tuple[int, int, int, int] | None = None, ): """Shows an image with optionally drawn in rectangle. Args: image: image read in with read_image_gray_scale title: title of the window rectangle: rectangle to draw on the image Returns: Opens a window to view the image. Close it with 'Space' to resume the scenario. """ image = image.copy() if rectangle is not None: cv2.rectangle(image, rectangle[0:2], rectangle[2:], (0, 255, 0), 2) cv2.imshow(title, image) cv2.waitKey() cv2.destroyAllWindows()
[docs] def read_image_gray_scale(image: bytes | os.PathLike | np.ndarray) -> cv2.Mat: """Read image in gray scale using opencv. Args: image: image object or path ot image Returns: the matrix containing the image """ if isinstance(image, bytes): target = np.frombuffer(image, np.uint8) target_gray = cv2.imdecode(target, cv2.IMREAD_GRAYSCALE) elif isinstance(image, (np.ndarray, cv2.Mat)): target = image.astype(np.uint8) if image.ndim == 2: return image target_gray = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY) else: target_gray = cv2.imread( str(pathlib.Path(image).resolve()), cv2.IMREAD_GRAYSCALE ) return target_gray
[docs] def read_image( image: bytes | os.PathLike | np.ndarray, color_code: int = cv2.IMREAD_COLOR, ) -> cv2.Mat: """Read an image with OpenCV Args: image: image to be read (can be bytes, path to image or numpy array) color_code: see OpenCV ImreadModes Returns: OpenCV image """ if isinstance(image, bytes): target = np.frombuffer(image, np.uint8) target = cv2.imdecode(target, color_code) elif isinstance(image, (np.ndarray, cv2.Mat)): target = image.astype(np.uint8) target = cv2.cvtColor(target, color_code) else: target = cv2.imread(str(pathlib.Path(image).resolve()), color_code) return target
def _resize_image( image_a: np.ndarray, image_b: np.ndarray ) -> tuple[np.ndarray, np.ndarray]: """Resize the bigger image to the smaller one. Do nothing if both images have the same shape. Args: image_a: first image to resize image_b: second image to resize Returns: Both images (one may be resized) in the same order, as they were provided to this function """ if image_a.shape == image_b.shape: return image_a, image_b elif image_a.shape < image_b.shape: return image_a, cv2.resize( image_b, dsize=image_a.shape[::-1], interpolation=cv2.INTER_CUBIC ) elif image_a.shape > image_b.shape: return ( cv2.resize( image_a, dsize=image_b.shape[::-1], interpolation=cv2.INTER_CUBIC ), image_b, )