zerohertzLib.vision.loader ¶

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[NDArray[np.uint8], list[int], list[NDArray[DTypeLike]]]:
    """
    Index에 따른 image와 `.txt` file에 대한 정보 return (`vis_path` 와 `class_color` 입력 시 시각화 image `vis_path` 에 저장)

    Args:
        idx: 입력 index

    Returns:
        읽어온 image와 그에 따른 `class_list` 및 `bbox` 혹은 `poly`
    """
    data_path = self.data_paths[idx]
    data_file_name = data_path.split("/")[-1]
    txt_path = os.path.join(
        self.txt_path, ".".join(data_file_name.split(".")[:-1]) + ".txt"
    )
    img = cv2.imread(data_path)
    try:
        class_list, objects = self._convert(txt_path, img)
    except FileNotFoundError:
        print(f"'{data_file_name}' is not found")
        return None, None, None
    if self.vis_path is not None:
        self._visualization(data_file_name, img, class_list, objects)
    return img, class_list, objects

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

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

def _annotation(self, args: list[int | str | list[str]]) -> dict[str, Any]:
    idx, directory, labels = args
    img, class_list, objects = self[idx]
    data_path = self.data_paths[idx]
    data_file_name = data_path.split("/")[-1]
    annotation = {
        "data": {"image": f"data/local-files/?d={directory}/{data_file_name}"}
    }
    result_data = []
    for cls, obj in zip(class_list, objects):
        result_data.append(self._value(img, obj, labels, cls))
    annotation["annotations"] = [{"result": result_data}]
    return annotation

def _convert(
    self, txt_path: str, img: NDArray[np.uint8]
) -> tuple[list[int], list[NDArray[DTypeLike]]]:
    class_list = []
    objects = []
    with open(txt_path, "r", encoding="utf-8") as file:
        data_lines = file.readlines()
    for data_line in data_lines:
        data_str = data_line.strip().split(" ")
        class_list.append(int(data_str[0]))
        if self.poly:
            obj = np.array(list(map(float, data_str[1:]))).reshape(-1, 2)
            if not self.absolute:
                obj *= img.shape[:2][::-1]
        else:
            obj = np.array(list(map(float, data_str[1:])))
            if not self.absolute:
                obj *= img.shape[:2][::-1] * 2
        objects.append(obj)
    return class_list, objects

def _value(
    self,
    img: NDArray[np.uint8],
    obj: NDArray[DTypeLike],
    labels: list[str],
    cls: int,
):
    original_height, original_width = img.shape[:2]
    obj *= 100
    if self.poly:
        obj /= (original_width, original_height)
        return {
            "original_width": original_width,
            "original_height": original_height,
            "image_rotation": 0,
            "value": {
                "points": obj.tolist(),
                "closed": True,
                "polygonlabels": [labels[cls]],
            },
            "from_name": "label",
            "to_name": "image",
            "type": "polygonlabels",
            "origin": "manual",
        }
    obj[:2] -= obj[2:] / 2
    obj /= (original_width, original_height) * 2
    obj = obj.tolist()
    return {
        "original_width": original_width,
        "original_height": original_height,
        "image_rotation": 0,
        "value": {
            "x": obj[0],
            "y": obj[1],
            "width": obj[2],
            "height": obj[3],
            "rectanglelabels": [labels[cls]],
        },
        "from_name": "label",
        "to_name": "image",
        "type": "rectanglelabels",
        "origin": "manual",
    }

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

def labelstudio(
    self,
    directory: str = "image",
    labels: list[str | None] = None,
    mp_num: int = 0,
) -> None:
    """
    YOLO format의 data를 Label Studio에서 확인 및 수정할 수 있게 변환

    Args:
        directory: Label Studio 내 `/home/user/{directory}` 의 이름
        labels: YOLO format의 `.txt` 상에서 index에 따른 label의 이름
        mp_num: 병렬 처리에 사용될 process의 수 (`0`: 직렬 처리)

    Returns:
        `{path}.json` 으로 결과 저장

    Examples
        >>> yolo.labelstudio("images", mp_num=10, labels=["t1", "t2", "t3", "t4"])
    """
    if labels is None:
        labels = [str(i) for i in range(100)]
    json_data = []
    if mp_num == 0:
        for idx in range(len(self)):
            json_data.append(self._annotation([idx, directory, labels]))
    else:
        args = [[idx, directory, labels] for idx in range(len(self))]
        with mp.Pool(processes=mp_num) as pool:
            annotations = pool.map(self._annotation, args)
        for annotation in annotations:
            json_data.append(annotation)
    write_json(json_data, self.data_path)