zerohertzLib.vision ¶

def __call__(
    self, idx: int, read: bool = False, int_class: bool = False
) -> tuple[
    str | NDArray[np.uint8],
    list[int | str],
    NDArray[DTypeLike],
    list[NDArray[DTypeLike]],
]:
    """
    Index에 따른 image와 annotation에 대한 정보 return (`vis_path` 와 `class_color` 입력 시 시각화 image `vis_path` 에 저장)

    Args:
        idx: 입력 index
        read: Image 읽음 여부
        int_class: 출력될 class의 type 지정

    Returns:
        Image 경로 혹은 읽어온 image와 그에 따른 `class_list`, `bboxes`, `polys`
    """
    img_path = os.path.join(
        self.data_path, os.path.basename(self.images[idx]["file_name"])
    )
    if read:
        img = cv2.imread(img_path)
    else:
        img = img_path
    class_list = []
    bboxes = []
    polys = []
    for idx_ in self.image2annotation[self.images[idx]["id"]]:
        annotation = self.annotations[idx_]
        if int_class:
            class_list.append(self.classes[annotation["category_id"]][0])
        else:
            class_list.append(self.classes[annotation["category_id"]][1])
        bboxes.append(
            [
                annotation["bbox"][0] + annotation["bbox"][2] / 2,
                annotation["bbox"][1] + annotation["bbox"][3] / 2,
                annotation["bbox"][2],
                annotation["bbox"][3],
            ]
        )
        if "segmentation" in annotation.keys():
            polys.append(np.array(annotation["segmentation"][0]).reshape(-1, 2))
    bboxes = np.array(bboxes)
    return img, class_list, bboxes, polys

def __getitem__(
    self, idx: int
) -> tuple[
    NDArray[np.uint8], list[str], NDArray[DTypeLike], list[NDArray[DTypeLike]]
]:
    """
    Index에 따른 image와 annotation에 대한 정보 return (`vis_path` 와 `class_color` 입력 시 시각화 image `vis_path` 에 저장)

    Args:
        idx: 입력 index

    Returns:
        읽어온 image와 그에 따른 `class_list`, `bboxes`, `polys`
    """
    img, class_list, bboxes, polys = self(idx, read=True)
    if self.vis_path is not None:
        self._visualization(
            os.path.basename(self.images[idx]["file_name"]),
            img,
            class_list,
            bboxes,
            polys,
        )
    return img, class_list, bboxes, polys

def __len__(self) -> int:
    """Image 수를 반환

    Returns:
        읽어온 image file들의 수
    """
    return len(self.images)

def _visualization(
    self,
    file_name: str,
    img: NDArray[np.uint8],
    class_list: list[str],
    bboxes: NDArray[DTypeLike],
    polys: list[NDArray[DTypeLike]],
) -> None:
    for cls, box in zip(class_list, bboxes):
        img = bbox(img, box, self.class_color[cls])
    if polys:
        mks = np.zeros((len(polys), *img.shape[:2]), bool)
        for idx, poly in enumerate(polys):
            mks[idx] = poly2mask(poly, img.shape[:2])
        img = mask(img, mks, class_list=class_list, class_color=self.class_color)
    cv2.imwrite(os.path.join(self.vis_path, file_name), img)

def yolo(
    self,
    target_path: str,
    label: list[str] | None = None,
    poly: bool = False,
) -> None:
    """COCO format을 YOLO format으로 변환

    Args:
        target_path: YOLO format data가 저장될 경로
        label: COCO에서 사용한 label을 정수로 변환하는 list (index 사용)
        poly: Segmentation format 유무

    Returns:
        `{target_path}/images` 및 `{target_path}/labels` 에 image와 `.txt` file 저장

    Examples:
        >>> coco = zz.vision.CocoLoader(data_path)
        >>> coco.yolo(target_path)
        100%|█████████████| 476/476 [00:00<00:00, 78794.25it/s]
        >>> label = ["label1", "label2"]
        >>> cooc.yolo(target_path, label)
        100%|█████████████| 476/476 [00:00<00:00, 78794.25it/s]
    """
    rmtree(os.path.join(target_path, "images"))
    rmtree(os.path.join(target_path, "labels"))
    for idx in tqdm(range(len(self))):
        img_path, class_list, bboxes, polys = self(
            idx, read=False, int_class=label is None
        )
        converted_gt = []
        if poly:
            for cls, poly_ in zip(class_list, polys):
                poly_ /= (self.images[idx]["width"], self.images[idx]["height"])
                if label:
                    cls = label.index(cls)
                converted_gt.append(
                    f"{cls} " + " ".join(map(str, poly_.reshape(-1)))
                )
        else:
            for cls, box in zip(class_list, bboxes):
                box /= (self.images[idx]["width"], self.images[idx]["height"]) * 2
                if label:
                    cls = label.index(cls)
                converted_gt.append(f"{cls} " + " ".join(map(str, box)))
        img_file_name = os.path.basename(img_path)
        txt_file_name = ".".join(img_file_name.split(".")[:-1]) + ".txt"
        try:
            shutil.copy(
                img_path, os.path.join(target_path, "images", img_file_name)
            )
            with open(
                os.path.join(target_path, "labels", txt_file_name),
                "w",
                encoding="utf-8",
            ) as file:
                file.writelines("\n".join(converted_gt))
        except FileNotFoundError:
            print(f"'{img_path}' is not found")

def __getitem__(
    self, idx: int
) -> tuple[str, NDArray[np.uint8]] | tuple[list[str], list[NDArray[np.uint8]]]:
    """Index에 따른 image 정보를 반환

    Args:
        idx: 입력 index

    Returns:
        `cnt` 에 따른 file 경로 및 image 값
    """
    if self.cnt == 1:
        return (
            self.image_paths[idx],
            cv2.imread(self.image_paths[idx], cv2.IMREAD_UNCHANGED),
        )
    return (
        self.image_paths[self.cnt * idx : self.cnt * (idx + 1)],
        [
            cv2.imread(path, cv2.IMREAD_UNCHANGED)
            for path in self.image_paths[self.cnt * idx : self.cnt * (idx + 1)]
        ],
    )

def __len__(self) -> int:
    """Image 수를 반환

    Returns:
        `cnt` 에 해당하는 image들의 수
    """
    return math.ceil(len(self.image_paths) / self.cnt)

def __getitem__(self, idx: int) -> tuple[NDArray[np.uint8], Json]:
    """
    읽어온 JSON file들을 list와 같이 indexing 후 해당하는 image return

    Args:
        idx: 입력 index

    Returns:
        Image와 JSON 내 정보
    """
    data_name = self.json[idx].get(self.json_key)
    img = cv2.imread(os.path.join(self.data_path, data_name), cv2.IMREAD_UNCHANGED)
    return img, self.json[idx]

def __len__(self) -> int:
    """Image 수를 반환

    Returns:
        읽어온 JSON file들의 수
    """
    return len(self.json)

def __getitem__(
    self, idx: int
) -> tuple[str, dict[str, dict[str, str]]] | tuple[str, dict[str, list[Any]]]:
    """
    Args:
        idx: 입력 index

    Returns:
        Index에 따른 image file 이름 또는 경로와 JSON file에 포함될 dictionary 또는 annotation 정보
    """
    if self.annotations is None:
        file_name = os.path.basename(self.data_paths[idx])
        return (
            file_name,
            {
                "data": {
                    "image": f"data/local-files/?d={self.path}/{self.data_paths[idx]}"
                }
            },
        )
    file_name = os.path.basename(self.annotations[idx]["data"]["image"])
    file_name = urllib.parse.unquote(file_name)
    if len(file_name) > 8 and "-" == file_name[8]:
        file_name = "-".join(file_name.split("-")[1:])
    file_path = os.path.join(self.data_path, file_name)
    if len(self.annotations[idx]["annotations"]) > 1:
        raise ValueError("The 'annotations' are plural")
    if self.type == "rectanglelabels":
        return (
            file_path,
            self._dict2cwh(self.annotations[idx]["annotations"][0]["result"]),
        )
    if self.type == "polygonlabels":
        return (
            file_path,
            self._dict2poly(self.annotations[idx]["annotations"][0]["result"]),
        )
    raise ValueError(f"Unknown annotation type: {self.type}")

def __len__(self) -> int:
    """데이터 개수를 반환

    Returns:
        읽어온 image file 혹은 annotation들의 수
    """
    if self.annotations is None:
        return len(self.data_paths)
    return len(self.annotations)

def _dict2cwh(self, results: list[dict[str, Any]]) -> dict[str, Any]:
    labels, polys, whs = [], [], []
    for result in results:
        width, height = result["original_width"], result["original_height"]
        box_cwh = (
            np.array(
                [
                    result["value"]["x"],
                    result["value"]["y"],
                    result["value"]["width"],
                    result["value"]["height"],
                ]
            )
            / 100
        )
        if len(result["value"]["rectanglelabels"]) > 1:
            raise ValueError("The 'rectanglelabels' are plural")
        label = result["value"]["rectanglelabels"][0]
        labels.append(label)
        self.labels.add(label)
        polys.append(box_cwh)
        whs.append((width, height))
    return {"labels": labels, "polys": polys, "whs": whs}

def _dict2poly(self, results: list[dict[str, Any]]) -> dict[str, Any]:
    labels, polys, whs = [], [], []
    for result in results:
        width, height = result["original_width"], result["original_height"]
        box_poly = np.array(result["value"]["points"]) / 100
        if len(result["value"]["polygonlabels"]) > 1:
            raise ValueError("The 'polygonlabels' are plural")
        label = result["value"]["polygonlabels"][0]
        labels.append(label)
        self.labels.add(label)
        polys.append(box_poly)
        whs.append((width, height))
    return {"labels": labels, "polys": polys, "whs": whs}

def classification(
    self,
    target_path: str,
    label: dict[str, Any] | None = None,
    rand: int = 0,
    shrink: bool = True,
    aug: int = 1,
) -> None:
    """Label Studio로 annotation한 JSON data를 classification format으로 변환

    Args:
        target_path: Classification format data가 저장될 경로
        label: Label Studio에서 사용한 label을 변경하는 dictionary
        rand: Image crop 시 random 범위 추가
        shrink: `rand` 에 의한 crop 시 image의 수축 여부
        aug: 한 annotation 당 저장할 image의 수

    Returns:
        annotation의 index, `i`: `rand` 의 index)

    Examples:
        >>> ls = zz.vision.LabelStudio(data_path, json_path)
        >>> ls.classification(target_path)
        100%|█████████████| 476/476 [00:00<00:00, 78794.25it/s]
        >>> label = {"label1": "lab1", "label2": "lab2"}
        >>> ls.classification(target_path, label, rand=10, aug=10, shrink=False)
        100%|█████████████| 476/476 [00:00<00:00, 78794.25it/s]
    """
    if label is None:
        label = {}
    for file_path, result in tqdm(self):
        img = cv2.imread(file_path)
        if img is None:
            print(f"'{file_path}' is not found")
            continue
        img_file = file_path.split("/")[-1].split(".")
        img_file_name = ".".join(img_file[:-1])
        img_file_ext = img_file[-1]
        for idx, (lab, poly, wh) in enumerate(
            zip(result["labels"], result["polys"], result["whs"])
        ):
            if self.type == "rectanglelabels":
                box_xyxy = poly * (wh * 2)
                box_xyxy[2:] += box_xyxy[:2]
            elif self.type == "polygonlabels":
                box_poly = poly * wh
                box_xyxy = poly2xyxy(box_poly)
            else:
                raise ValueError(f"Unknown annotation type: {self.type}")
            os.makedirs(
                os.path.join(target_path, label.get(lab, lab)), exist_ok=True
            )
            for i in range(aug):
                bias = (2 * rand * (np.random.rand(4) - 0.5)).astype(np.int32)
                if not shrink:
                    bias[:2] = -abs(bias[:2])
                    bias[2:] = abs(bias[2:])
                x_0, y_0, x_1, y_1 = box_xyxy.astype(np.int32) + bias
                try:
                    cv2.imwrite(
                        os.path.join(
                            target_path,
                            label.get(lab, lab),
                            f"{img_file_name}_{idx}_{i}.{img_file_ext}",
                        ),
                        img[y_0:y_1, x_0:x_1, :],
                    )
                except cv2.error:
                    print(
                        f"Impossible crop ('x_0': {x_0}, 'y_0': {y_0}, 'x_1': {x_1}, 'y_1': {y_1})"
                    )

def coco(self, target_path: str, label: dict[str, int]) -> None:
    """Label Studio로 annotation한 JSON data를 COCO format으로 변환

    Args:
        target_path: COCO format data가 저장될 경로
        label: Label Studio에서 사용한 label을 변경하는 dictionary

    Returns:
        `{target_path}.json` 에 JSON file 저장

    Examples:
        >>> ls = zz.vision.LabelStudio(data_path, json_path)
        >>> label = {"label1": 1, "label2": 2}
        >>> ls.coco(target_path, label)
        100%|█████████████| 476/476 [00:00<00:00, 78794.25it/s]
    """
    converted_gt = {
        "images": [],
        "annotations": [],
        "categories": [],
    }
    for lab, id_ in label.items():
        converted_gt["categories"].append({"id": id_, "name": lab})
    ant_id = 0
    for id_, (file_path, result) in enumerate(tqdm(self)):
        _images = {
            "file_name": os.path.basename(file_path),
            "height": result["whs"][0][1],
            "width": result["whs"][0][0],
            "id": id_,
        }
        _annotations = []
        for ant_id_, (lab, poly, wh) in enumerate(
            zip(result["labels"], result["polys"], result["whs"])
        ):
            # box_cwh is [x_0, y_0, width, height] not [cx, cy, width, height]
            if self.type == "rectanglelabels":
                poly = poly * (wh * 2)
                box_cwh = poly.copy()
                poly[2:] += poly[:2]
                poly = xyxy2poly(poly)
            elif self.type == "polygonlabels":
                poly = poly * wh
                box_cwh = poly2cwh(poly)
                box_cwh[:2] -= box_cwh[2:] / 2
            else:
                raise ValueError(f"Unknown annotation type: {self.type}")
            box_cwh = box_cwh.tolist()
            _annotations.append(
                {
                    "segmentation": [poly.reshape(-1).tolist()],
                    "area": box_cwh[2] * box_cwh[3],
                    "iscrowd": 0,
                    "image_id": id_,
                    "bbox": box_cwh,
                    "category_id": label[lab],
                    "id": ant_id + ant_id_,
                }
            )
        converted_gt["images"].append(_images)
        converted_gt["annotations"] += _annotations
        ant_id += len(result["labels"]) + 1
    write_json(converted_gt, target_path)