bplusplus 1.1.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

bplusplus/yolov5detect/utils/segment/dataloaders.py

@@ -1,366 +0,0 @@
-# Ultralytics YOLOv5 🚀, AGPL-3.0 license
-"""Dataloaders."""
-
-import os
-import random
-
-import cv2
-import numpy as np
-import torch
-from torch.utils.data import DataLoader
-
-from ..augmentations import augment_hsv, copy_paste, letterbox
-from ..dataloaders import InfiniteDataLoader, LoadImagesAndLabels, SmartDistributedSampler, seed_worker
-from ..general import LOGGER, xyn2xy, xywhn2xyxy, xyxy2xywhn
-from ..torch_utils import torch_distributed_zero_first
-from .augmentations import mixup, random_perspective
-
-RANK = int(os.getenv("RANK", -1))
-
-
-def create_dataloader(
-    path,
-    imgsz,
-    batch_size,
-    stride,
-    single_cls=False,
-    hyp=None,
-    augment=False,
-    cache=False,
-    pad=0.0,
-    rect=False,
-    rank=-1,
-    workers=8,
-    image_weights=False,
-    quad=False,
-    prefix="",
-    shuffle=False,
-    mask_downsample_ratio=1,
-    overlap_mask=False,
-    seed=0,
-):
-    """Creates a dataloader for training, validating, or testing YOLO models with various dataset options."""
-    if rect and shuffle:
-        LOGGER.warning("WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False")
-        shuffle = False
-    with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
-        dataset = LoadImagesAndLabelsAndMasks(
-            path,
-            imgsz,
-            batch_size,
-            augment=augment,  # augmentation
-            hyp=hyp,  # hyperparameters
-            rect=rect,  # rectangular batches
-            cache_images=cache,
-            single_cls=single_cls,
-            stride=int(stride),
-            pad=pad,
-            image_weights=image_weights,
-            prefix=prefix,
-            downsample_ratio=mask_downsample_ratio,
-            overlap=overlap_mask,
-            rank=rank,
-        )
-
-    batch_size = min(batch_size, len(dataset))
-    nd = torch.cuda.device_count()  # number of CUDA devices
-    nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers])  # number of workers
-    sampler = None if rank == -1 else SmartDistributedSampler(dataset, shuffle=shuffle)
-    loader = DataLoader if image_weights else InfiniteDataLoader  # only DataLoader allows for attribute updates
-    generator = torch.Generator()
-    generator.manual_seed(6148914691236517205 + seed + RANK)
-    return loader(
-        dataset,
-        batch_size=batch_size,
-        shuffle=shuffle and sampler is None,
-        num_workers=nw,
-        sampler=sampler,
-        pin_memory=True,
-        collate_fn=LoadImagesAndLabelsAndMasks.collate_fn4 if quad else LoadImagesAndLabelsAndMasks.collate_fn,
-        worker_init_fn=seed_worker,
-        generator=generator,
-    ), dataset
-
-
-class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
-    """Loads images, labels, and segmentation masks for training and testing YOLO models with augmentation support."""
-
-    def __init__(
-        self,
-        path,
-        img_size=640,
-        batch_size=16,
-        augment=False,
-        hyp=None,
-        rect=False,
-        image_weights=False,
-        cache_images=False,
-        single_cls=False,
-        stride=32,
-        pad=0,
-        min_items=0,
-        prefix="",
-        downsample_ratio=1,
-        overlap=False,
-        rank=-1,
-        seed=0,
-    ):
-        """Initializes the dataset with image, label, and mask loading capabilities for training/testing."""
-        super().__init__(
-            path,
-            img_size,
-            batch_size,
-            augment,
-            hyp,
-            rect,
-            image_weights,
-            cache_images,
-            single_cls,
-            stride,
-            pad,
-            min_items,
-            prefix,
-            rank,
-            seed,
-        )
-        self.downsample_ratio = downsample_ratio
-        self.overlap = overlap
-
-    def __getitem__(self, index):
-        """Returns a transformed item from the dataset at the specified index, handling indexing and image weighting."""
-        index = self.indices[index]  # linear, shuffled, or image_weights
-
-        hyp = self.hyp
-        mosaic = self.mosaic and random.random() < hyp["mosaic"]
-        masks = []
-        if mosaic:
-            # Load mosaic
-            img, labels, segments = self.load_mosaic(index)
-            shapes = None
-
-            # MixUp augmentation
-            if random.random() < hyp["mixup"]:
-                img, labels, segments = mixup(img, labels, segments, *self.load_mosaic(random.randint(0, self.n - 1)))
-
-        else:
-            # Load image
-            img, (h0, w0), (h, w) = self.load_image(index)
-
-            # Letterbox
-            shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size  # final letterboxed shape
-            img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
-            shapes = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling
-
-            labels = self.labels[index].copy()
-            # [array, array, ....], array.shape=(num_points, 2), xyxyxyxy
-            segments = self.segments[index].copy()
-            if len(segments):
-                for i_s in range(len(segments)):
-                    segments[i_s] = xyn2xy(
-                        segments[i_s],
-                        ratio[0] * w,
-                        ratio[1] * h,
-                        padw=pad[0],
-                        padh=pad[1],
-                    )
-            if labels.size:  # normalized xywh to pixel xyxy format
-                labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
-
-            if self.augment:
-                img, labels, segments = random_perspective(
-                    img,
-                    labels,
-                    segments=segments,
-                    degrees=hyp["degrees"],
-                    translate=hyp["translate"],
-                    scale=hyp["scale"],
-                    shear=hyp["shear"],
-                    perspective=hyp["perspective"],
-                )
-
-        nl = len(labels)  # number of labels
-        if nl:
-            labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3)
-            if self.overlap:
-                masks, sorted_idx = polygons2masks_overlap(
-                    img.shape[:2], segments, downsample_ratio=self.downsample_ratio
-                )
-                masks = masks[None]  # (640, 640) -> (1, 640, 640)
-                labels = labels[sorted_idx]
-            else:
-                masks = polygons2masks(img.shape[:2], segments, color=1, downsample_ratio=self.downsample_ratio)
-
-        masks = (
-            torch.from_numpy(masks)
-            if len(masks)
-            else torch.zeros(
-                1 if self.overlap else nl, img.shape[0] // self.downsample_ratio, img.shape[1] // self.downsample_ratio
-            )
-        )
-        # TODO: albumentations support
-        if self.augment:
-            # Albumentations
-            # there are some augmentation that won't change boxes and masks,
-            # so just be it for now.
-            img, labels = self.albumentations(img, labels)
-            nl = len(labels)  # update after albumentations
-
-            # HSV color-space
-            augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"])
-
-            # Flip up-down
-            if random.random() < hyp["flipud"]:
-                img = np.flipud(img)
-                if nl:
-                    labels[:, 2] = 1 - labels[:, 2]
-                    masks = torch.flip(masks, dims=[1])
-
-            # Flip left-right
-            if random.random() < hyp["fliplr"]:
-                img = np.fliplr(img)
-                if nl:
-                    labels[:, 1] = 1 - labels[:, 1]
-                    masks = torch.flip(masks, dims=[2])
-
-            # Cutouts  # labels = cutout(img, labels, p=0.5)
-
-        labels_out = torch.zeros((nl, 6))
-        if nl:
-            labels_out[:, 1:] = torch.from_numpy(labels)
-
-        # Convert
-        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
-        img = np.ascontiguousarray(img)
-
-        return (torch.from_numpy(img), labels_out, self.im_files[index], shapes, masks)
-
-    def load_mosaic(self, index):
-        """Loads 1 image + 3 random images into a 4-image YOLOv5 mosaic, adjusting labels and segments accordingly."""
-        labels4, segments4 = [], []
-        s = self.img_size
-        yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border)  # mosaic center x, y
-
-        # 3 additional image indices
-        indices = [index] + random.choices(self.indices, k=3)  # 3 additional image indices
-        for i, index in enumerate(indices):
-            # Load image
-            img, _, (h, w) = self.load_image(index)
-
-            # place img in img4
-            if i == 0:  # top left
-                img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
-                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
-                x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)
-            elif i == 1:  # top right
-                x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
-                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
-            elif i == 2:  # bottom left
-                x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
-                x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
-            elif i == 3:  # bottom right
-                x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
-                x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
-
-            img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
-            padw = x1a - x1b
-            padh = y1a - y1b
-
-            labels, segments = self.labels[index].copy(), self.segments[index].copy()
-
-            if labels.size:
-                labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh)  # normalized xywh to pixel xyxy format
-                segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
-            labels4.append(labels)
-            segments4.extend(segments)
-
-        # Concat/clip labels
-        labels4 = np.concatenate(labels4, 0)
-        for x in (labels4[:, 1:], *segments4):
-            np.clip(x, 0, 2 * s, out=x)  # clip when using random_perspective()
-        # img4, labels4 = replicate(img4, labels4)  # replicate
-
-        # Augment
-        img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"])
-        img4, labels4, segments4 = random_perspective(
-            img4,
-            labels4,
-            segments4,
-            degrees=self.hyp["degrees"],
-            translate=self.hyp["translate"],
-            scale=self.hyp["scale"],
-            shear=self.hyp["shear"],
-            perspective=self.hyp["perspective"],
-            border=self.mosaic_border,
-        )  # border to remove
-        return img4, labels4, segments4
-
-    @staticmethod
-    def collate_fn(batch):
-        """Custom collation function for DataLoader, batches images, labels, paths, shapes, and segmentation masks."""
-        img, label, path, shapes, masks = zip(*batch)  # transposed
-        batched_masks = torch.cat(masks, 0)
-        for i, l in enumerate(label):
-            l[:, 0] = i  # add target image index for build_targets()
-        return torch.stack(img, 0), torch.cat(label, 0), path, shapes, batched_masks
-
-
-def polygon2mask(img_size, polygons, color=1, downsample_ratio=1):
-    """
-    Args:
-        img_size (tuple): The image size.
-        polygons (np.ndarray): [N, M], N is the number of polygons,
-            M is the number of points(Be divided by 2).
-    """
-    mask = np.zeros(img_size, dtype=np.uint8)
-    polygons = np.asarray(polygons)
-    polygons = polygons.astype(np.int32)
-    shape = polygons.shape
-    polygons = polygons.reshape(shape[0], -1, 2)
-    cv2.fillPoly(mask, polygons, color=color)
-    nh, nw = (img_size[0] // downsample_ratio, img_size[1] // downsample_ratio)
-    # NOTE: fillPoly firstly then resize is trying the keep the same way
-    # of loss calculation when mask-ratio=1.
-    mask = cv2.resize(mask, (nw, nh))
-    return mask
-
-
-def polygons2masks(img_size, polygons, color, downsample_ratio=1):
-    """
-    Args:
-        img_size (tuple): The image size.
-        polygons (list[np.ndarray]): each polygon is [N, M],
-            N is the number of polygons,
-            M is the number of points(Be divided by 2).
-    """
-    masks = []
-    for si in range(len(polygons)):
-        mask = polygon2mask(img_size, [polygons[si].reshape(-1)], color, downsample_ratio)
-        masks.append(mask)
-    return np.array(masks)
-
-
-def polygons2masks_overlap(img_size, segments, downsample_ratio=1):
-    """Return a (640, 640) overlap mask."""
-    masks = np.zeros(
-        (img_size[0] // downsample_ratio, img_size[1] // downsample_ratio),
-        dtype=np.int32 if len(segments) > 255 else np.uint8,
-    )
-    areas = []
-    ms = []
-    for si in range(len(segments)):
-        mask = polygon2mask(
-            img_size,
-            [segments[si].reshape(-1)],
-            downsample_ratio=downsample_ratio,
-            color=1,
-        )
-        ms.append(mask)
-        areas.append(mask.sum())
-    areas = np.asarray(areas)
-    index = np.argsort(-areas)
-    ms = np.array(ms)[index]
-    for i in range(len(segments)):
-        mask = ms[i] * (i + 1)
-        masks = masks + mask
-        masks = np.clip(masks, a_min=0, a_max=i + 1)
-    return masks, index

bplusplus/yolov5detect/utils/segment/general.py

@@ -1,160 +0,0 @@
-# Ultralytics YOLOv5 🚀, AGPL-3.0 license
-
-import cv2
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-
-def crop_mask(masks, boxes):
-    """
-    "Crop" predicted masks by zeroing out everything not in the predicted bbox. Vectorized by Chong (thanks Chong).
-
-    Args:
-        - masks should be a size [n, h, w] tensor of masks
-        - boxes should be a size [n, 4] tensor of bbox coords in relative point form
-    """
-    n, h, w = masks.shape
-    x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1)  # x1 shape(1,1,n)
-    r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :]  # rows shape(1,w,1)
-    c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None]  # cols shape(h,1,1)
-
-    return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
-
-
-def process_mask_upsample(protos, masks_in, bboxes, shape):
-    """
-    Crop after upsample.
-    protos: [mask_dim, mask_h, mask_w]
-    masks_in: [n, mask_dim], n is number of masks after nms
-    bboxes: [n, 4], n is number of masks after nms
-    shape: input_image_size, (h, w).
-
-    return: h, w, n
-    """
-    c, mh, mw = protos.shape  # CHW
-    masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
-    masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
-    masks = crop_mask(masks, bboxes)  # CHW
-    return masks.gt_(0.5)
-
-
-def process_mask(protos, masks_in, bboxes, shape, upsample=False):
-    """
-    Crop before upsample.
-    proto_out: [mask_dim, mask_h, mask_w]
-    out_masks: [n, mask_dim], n is number of masks after nms
-    bboxes: [n, 4], n is number of masks after nms
-    shape:input_image_size, (h, w).
-
-    return: h, w, n
-    """
-    c, mh, mw = protos.shape  # CHW
-    ih, iw = shape
-    masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)  # CHW
-
-    downsampled_bboxes = bboxes.clone()
-    downsampled_bboxes[:, 0] *= mw / iw
-    downsampled_bboxes[:, 2] *= mw / iw
-    downsampled_bboxes[:, 3] *= mh / ih
-    downsampled_bboxes[:, 1] *= mh / ih
-
-    masks = crop_mask(masks, downsampled_bboxes)  # CHW
-    if upsample:
-        masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
-    return masks.gt_(0.5)
-
-
-def process_mask_native(protos, masks_in, bboxes, shape):
-    """
-    Crop after upsample.
-    protos: [mask_dim, mask_h, mask_w]
-    masks_in: [n, mask_dim], n is number of masks after nms
-    bboxes: [n, 4], n is number of masks after nms
-    shape: input_image_size, (h, w).
-
-    return: h, w, n
-    """
-    c, mh, mw = protos.shape  # CHW
-    masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
-    gain = min(mh / shape[0], mw / shape[1])  # gain  = old / new
-    pad = (mw - shape[1] * gain) / 2, (mh - shape[0] * gain) / 2  # wh padding
-    top, left = int(pad[1]), int(pad[0])  # y, x
-    bottom, right = int(mh - pad[1]), int(mw - pad[0])
-    masks = masks[:, top:bottom, left:right]
-
-    masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
-    masks = crop_mask(masks, bboxes)  # CHW
-    return masks.gt_(0.5)
-
-
-def scale_image(im1_shape, masks, im0_shape, ratio_pad=None):
-    """
-    img1_shape: model input shape, [h, w]
-    img0_shape: origin pic shape, [h, w, 3]
-    masks: [h, w, num].
-    """
-    # Rescale coordinates (xyxy) from im1_shape to im0_shape
-    if ratio_pad is None:  # calculate from im0_shape
-        gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1])  # gain  = old / new
-        pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2  # wh padding
-    else:
-        pad = ratio_pad[1]
-    top, left = int(pad[1]), int(pad[0])  # y, x
-    bottom, right = int(im1_shape[0] - pad[1]), int(im1_shape[1] - pad[0])
-
-    if len(masks.shape) < 2:
-        raise ValueError(f'"len of masks shape" should be 2 or 3, but got {len(masks.shape)}')
-    masks = masks[top:bottom, left:right]
-    # masks = masks.permute(2, 0, 1).contiguous()
-    # masks = F.interpolate(masks[None], im0_shape[:2], mode='bilinear', align_corners=False)[0]
-    # masks = masks.permute(1, 2, 0).contiguous()
-    masks = cv2.resize(masks, (im0_shape[1], im0_shape[0]))
-
-    if len(masks.shape) == 2:
-        masks = masks[:, :, None]
-    return masks
-
-
-def mask_iou(mask1, mask2, eps=1e-7):
-    """
-    mask1: [N, n] m1 means number of predicted objects
-    mask2: [M, n] m2 means number of gt objects
-    Note: n means image_w x image_h.
-
-    return: masks iou, [N, M]
-    """
-    intersection = torch.matmul(mask1, mask2.t()).clamp(0)
-    union = (mask1.sum(1)[:, None] + mask2.sum(1)[None]) - intersection  # (area1 + area2) - intersection
-    return intersection / (union + eps)
-
-
-def masks_iou(mask1, mask2, eps=1e-7):
-    """
-    mask1: [N, n] m1 means number of predicted objects
-    mask2: [N, n] m2 means number of gt objects
-    Note: n means image_w x image_h.
-
-    return: masks iou, (N, )
-    """
-    intersection = (mask1 * mask2).sum(1).clamp(0)  # (N, )
-    union = (mask1.sum(1) + mask2.sum(1))[None] - intersection  # (area1 + area2) - intersection
-    return intersection / (union + eps)
-
-
-def masks2segments(masks, strategy="largest"):
-    """Converts binary (n,160,160) masks to polygon segments with options for concatenation or selecting the largest
-    segment.
-    """
-    segments = []
-    for x in masks.int().cpu().numpy().astype("uint8"):
-        c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
-        if c:
-            if strategy == "concat":  # concatenate all segments
-                c = np.concatenate([x.reshape(-1, 2) for x in c])
-            elif strategy == "largest":  # select largest segment
-                c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
-        else:
-            c = np.zeros((0, 2))  # no segments found
-        segments.append(c.astype("float32"))
-    return segments