Spaces:
Running
Running
File size: 1,686 Bytes
b683920 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
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
|