water_turbidity_det/labelme/utils.py

import numpy as np
import cv2

def draw_grid(img: np.ndarray, grid_w: int, grid_h:int):
    """划格网"""
    img_h, img_w, _ = img.shape
    # 绘制横向网格线
    for i in range((img_h // grid_h)+1):
        cv2.line(img, (0, i*grid_h), (img_w, i*grid_h), (0, 255, 0), 2)
    # 绘制纵向网格线
    for i in range((img_w // grid_w)+1):
        cv2.line(img, (i*grid_w, 0), (i*grid_w, img_h), (0, 255, 0), 2)
    return img

def draw_predict_grid(img, patches_index, predicted_class, confidence):
    """在img上绘制预测结果"""
    for i, (idx_w, idx_h) in enumerate(patches_index):
        cv2.circle(img, (idx_w, idx_h), 10, (0, 255, 0), -1)
        text1 = f'cls:{predicted_class[i]}'
        text2 = f'prob:{confidence[i] * 100:.1f}%'
        color = (0, 0, 255) if predicted_class[i] else (255, 0, 0)
        cv2.putText(img, text1, (idx_w, idx_h + 128),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2)
        cv2.putText(img, text2, (idx_w, idx_h + 128 + 25),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2)
    return img