jsw/utils.py

import random
import cv2
import matplotlib.pyplot as plt


def draw_points(image, points, color=None, random_color=False, same=True, thickness=1):
    if color is None and not random_color:
        color = (0, 255, 0)  # Màu mặc định là xanh lá cây (BGR)
    if random_color:
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

    image = to_color(image)

    for point in points:
        if random_color and not same:
            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

        x, y = point
        image = cv2.circle(image, (x, y), thickness, color, -1)  # Vẽ điểm lên ảnh
    return image


def draw_lines(image, pairs, color=None, random_color=False, same=True, thickness=1):
    image_with_line = to_color(np.copy(image))

    if color is None and not random_color:
        color = (255, 0, 0)  # Màu mặc định là xanh lá cây (BGR)
    if random_color:
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

    # Vẽ đường thẳng dựa trên danh sách các cặp điểm
    for pair in pairs:

        if random_color and not same:
            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

        start_point = pair[0]
        end_point = pair[1]
        image_with_line = cv2.line(image_with_line, start_point, end_point, color, thickness)
        image_with_line = cv2.circle(image_with_line, start_point, thickness + 1, color, -1)
        image_with_line = cv2.circle(image_with_line, end_point, thickness + 1, color, -1)

    return image_with_line


def center(contour):
    idx = len(contour) // 2
    return contour[idx]


def to_color(image):
    if len(image.shape) == 3 and image.shape[-1] == 3:
        return image
    return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)


def to_gray(image):
    if len(image.shape) == 3 and image.shape[-1] == 3:
        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return image


def cv2_imshow(images):
    if not isinstance(images, list):
        images = [images]

    num_images = len(images)

    # Hiển thị ảnh đơn lẻ trực tiếp bằng imshow
    if num_images == 1:
        image = images[0]
        if len(image.shape) == 3 and image.shape[2] == 3:
            # Ảnh màu (RGB)
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            plt.imshow(image_rgb)
        else:
            # Ảnh xám
            plt.imshow(image, cmap='gray')

        plt.axis("off")
        plt.show()
    else:
        # Hiển thị nhiều ảnh trên cùng một cột
        fig, ax = plt.subplots(num_images, 1, figsize=(4, 4 * num_images))

        for i in range(num_images):
            image = images[i]
            if len(image.shape) == 3 and image.shape[2] == 3:
                # Ảnh màu (RGB)
                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                ax[i].imshow(image_rgb)
            else:
                # Ảnh xám
                ax[i].imshow(image, cmap='gray')

            ax[i].axis("off")

        plt.tight_layout()
        plt.show()