cvtk.ml.data

class cvtk.ml.data.DataLabel(labels: list | tuple | str)

Class to treat class labels

This class is designed to manage class (category) labels for machine learning tasks. The class loads class labels from a list, tuple, or text file when creating an instance. Methods implemented in the class provide a way to get the class index from the class name and vice versa.

Parameters:

labels – A tuple or list, or a path to a text file or coco format containing class labels. Text file should contain one class name per line.

Examples

>>> from cvtk.ml import DataLabel
>>>
>>> labels = ['leaf', 'flower', 'root']
>>> DataLabel = DataLabel(labels)
>>> print(DataLabel[1])
'flower'
>>> print(DataLabel['flower'])
1
>>> len(DataLabel)
3
>>> DataLabel.classes
['leaf', 'flower', 'root']
>>>
>>>
>>> labels = 'labels.txt'
>>> DataLabel = DataLabel(labels)
>>> print(DataLabel[1])
'flower'
>>> print(DataLabel['flower'])
1

class cvtk.ml.data.SquareResize(shape: int = 600, bg_color: tuple[int, int, int] | None = None, resample: object = 2)

Resize an image to a square shape

SquareResize provides a function to resize an image to a square. The resizing process resizes the length of the long side of the input image to the specified size, then adds padding to both sides of the short side of the image to convert the image to a square.

The background of the padding area is set by default to stretch the pixels at both ends of the image as is and then blur them. By specifying bg_color, the background of the padding area can be set to a single color.

Parameters:

• shape (int) – The resolution of the square image.

• bg_color (tuple) – The color of the padding area. Default is None. If None, the color is extended from both ends of the image.

• resample (int) – The resampling filter to use. Default is PIL.Image.BILINEAR.

Returns:

The squre image in PIL.Image.Image class.

Examples

>>> from cvtk.ml import SquareResize
>>>
>>> squareresize = SquareResize(shape=600)
>>> img = squareresize('image.jpg')
>>> img.save('image_square.jpg')
>>>
>>> squareresize = SquareResize(shape=600, bg_color=(0, 0, 0))
>>> img = squareresize('image.jpg')
>>> img.save('image_square.jpg')
>>>
>>>
>>> import torchvision.transforms
>>> transform = torchvision.transforms.Compose([
        SquareResize(256),
        torchvision.transforms.RandomHorizontalFlip(0.5),
        torchvision.transforms.RandomAffine(45),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.485, 0.456, 0.406],
                                         [0.229, 0.224, 0.225])
    ])