Skip to content

saige_vision_reader

data.dataset.saige_vision_reader

load_labels 

load_labels(file: Dict) -> List[Dict]
Source code in SaigeToolkit/data/dataset/saige_vision_reader.py 
12
13
14
15
16
17
18
19
20
21
def load_labels(file: Dict) -> List[Dict]:
    if "labels" in file:
        labels = file["labels"]
    elif "labels_path" in file:
        with open(file["labels_path"], "rb") as f:
            labels = pickle.load(f)
    else:
        labels = []

    return labels

load_label_cls 

load_label_cls(file: Dict) -> Dict
Source code in SaigeToolkit/data/dataset/saige_vision_reader.py 
24
25
def load_label_cls(file: Dict) -> Dict:
    return {"class": file["class_index"]}

load_label_iad 

load_label_iad(file: Dict) -> Dict
Source code in SaigeToolkit/data/dataset/saige_vision_reader.py 
28
29
def load_label_iad(file: Dict) -> Dict:
    return {"is_ng": file["labels"]["is_ng"]}

load_label_det 

load_label_det(file: Dict) -> Dict

박스 별 dict를 concat 되어있는 bboxes, labels, scores로 변환합니다.

Source code in SaigeToolkit/data/dataset/saige_vision_reader.py 
32
33
34
35
36
37
38
39
40
41
42
43
def load_label_det(file: Dict) -> Dict:
    """박스 별 dict를 concat 되어있는 bboxes, labels, scores로 변환합니다."""
    bboxes = file["labels"]

    label = {
        "bboxes": [bbox["bounding_box"] for bbox in bboxes],
        "labels": [bbox["class_index"] for bbox in bboxes],
    }
    if bboxes and "score" in bboxes[0]:
        label["scores"] = [bbox["score"] for bbox in bboxes]

    return label

load_label_seg 

load_label_seg(file: Dict, get_objects: bool = False, n_classes: Optional[int] = None) -> Dict
Source code in SaigeToolkit/data/dataset/saige_vision_reader.py 
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def load_label_seg(file: Dict, get_objects: bool = False, n_classes: Optional[int] = None) -> Dict:
    if "labels" in file:
        labels = file["labels"]
    elif "labels_path" in file:
        with open(file["labels_path"], "rb") as f:
            labels = pickle.load(f)
    else:  # unlabeled
        return {}

    w, h = file["width"], file["height"]
    mask = np.zeros((h, w), dtype=np.uint8)

    if get_objects:
        objects = []

    for label in labels:
        class_index = label["class_index"]
        if class_index == 0:
            raise ValueError("segment class_index cannot be 0")
        elif n_classes is not None and class_index >= n_classes:
            raise ValueError(
                f"class_index must be smaller than n_classes: got class_index={class_index} & n_classes={n_classes}"
            )

        bitmap_nonzero = None

        if "contours" in label:
            if not isinstance(label["contours"][0], np.ndarray):
                # NOTE: this is inplace
                label["contours"] = [np.array(contour) for contour in label["contours"]]
            bbox, bitmap = convert_contours_to_box_and_bitmap(label["contours"])
            bbox_xywh = [bbox.left, bbox.top, bbox.width, bbox.height]
        else:
            bbox_xywh = label["bounding_box"]
            bitmap = label["bitmap"]
            if not isinstance(bitmap, np.ndarray):
                bitmap = Image.open(BytesIO(base64.b64decode(bitmap)))
                bitmap = np.array(bitmap)
            bitmap_nonzero = bitmap > 0
            bitmap[bitmap_nonzero] = BITMAP_PIXEL_VALUE  # NOTE: this is inplace

        bbox_x, bbox_y, bbox_w, bbox_h = bbox_xywh
        bbox_bottom = bbox_y + bbox_h
        bbox_right = bbox_x + bbox_w
        if bbox_x < 0 or bbox_y < 0 or bbox_bottom > h or bbox_right > w:
            raise ValueError(f"segment is out of image: image_size={(w,h)}, bbox={(bbox_xywh)}")
        mask_sliced = mask[bbox_y:bbox_bottom, bbox_x:bbox_right]
        if bitmap_nonzero is None:
            bitmap_nonzero = bitmap > 0
        mask_sliced[bitmap_nonzero] = class_index

        if get_objects:
            segment = {
                "bounding_box": bbox_xywh,
                "bitmap": bitmap,
                "class_index": label["class_index"],
            }
            if "contours" in label:
                segment["contours"] = label["contours"]
            objects.append(segment)

    label = {"mask": mask}
    if get_objects:
        label["objects"] = objects

    return label

load_label_ocr 

load_label_ocr(file: Dict) -> Dict
Source code in SaigeToolkit/data/dataset/saige_vision_reader.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def load_label_ocr(file: Dict) -> Dict:
    labels = load_labels(file)

    pols, strs, kies, verts = [], [], [], []
    for label in labels:
        pols.append(np.array(label["polygon"]))
        strs.append(label["text"])
        if "class_index" in label:
            kies.append(label["class_index"])
        if "is_vertical" in label:
            verts.append(label["is_vertical"])

    ignore = [s is None for s in strs]

    label = {"polygons": pols, "strings": strs, "ignore": ignore}
    if len(kies) == len(pols):
        label.update({"kie_labels": kies})
    if len(verts) == len(pols):
        label.update({"is_vertical": verts})

    return label