from typing import List, Dict import matplotlib import matplotlib.patches as patches import matplotlib.pyplot as plt import numpy as np import tomli as tomllib def plot_img_with_rects( img: np.array, boxes: List[List], threshold: float = 0.5, coef: int = 400 ) -> matplotlib.figure.Figure: """ Plot image with rectangles. Args: img: image as a numpy array boxes: the list of the bboxes threshold: threshold for bbox probability coef: coefficient to multiply images. Can be changed when the original image is a different size Returns: image with bboxes """ fig, ax = plt.subplots(1, figsize=(4, 4)) # Display the image ax.imshow(img) # Create a Rectangle patch for _, rect in enumerate([b for b in boxes if b[1] > threshold]): label, _, xc, yc, w, h = rect xc, yc, w, h = xc * coef, yc * coef, w * coef, h * coef # the coordinates from center-based to left top corner x = xc - w / 2 y = yc - h / 2 label = int(label) label = label if label != 10 else 'penis' label = label if label != 11 else 'junk' rect = [x, y, x + w, y + h] rect_ = patches.Rectangle( (rect[0], rect[1]), rect[2] - rect[0], rect[3] - rect[1], linewidth=2, edgecolor='blue', facecolor='none' ) plt.text(rect[2], rect[1], f'{label}', color='blue') # Add the patch to the Axes ax.add_patch(rect_) return fig def get_config() -> Dict: """ Get dict from config. Returns: config """ with open('config.toml', 'rb') as f: config = tomllib.load(f) return config