uidetection/separate_datasets.py at main · RPA-US/uidetection · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
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
112
113
114
115
116
117
118
119
120
121
122
123
import os
import json

from shutil import copyfile
from shapely.geometry import Polygon

prefix = "Datasets/Dataset"

if not os.path.exists(f"{prefix}_TopLevel"):
    os.mkdir(f"{prefix}_TopLevel")

if not os.path.exists(f"{prefix}_ApplicationLevel"):
    os.mkdir(f"{prefix}_ApplicationLevel")

if not os.path.exists(f"{prefix}_ContainerLevel"):
    os.mkdir(f"{prefix}_ContainerLevel")

if not os.path.exists(f"{prefix}_ElementLevel"):
    os.mkdir(f"{prefix}_ElementLevel")

if not os.path.exists(f"{prefix}_TextLevel"):
    os.mkdir(f"{prefix}_TextLevel")

for file in os.listdir(prefix):
    if file.endswith(".json"):
        labeled_json = json.load(open(f"{prefix}/{file}"))

        top_level_json = labeled_json.copy()
        top_level_json["shapes"] = list(
            filter(
                lambda x: x["label"] in ["Application", "Taskbar", "Dock"],
                top_level_json["shapes"],
            )
        )

        application_level_json = labeled_json.copy()
        application_level_json["shapes"] = list(
            filter(
                lambda x: x["label"]
                in ["Header", "BrowserToolbar", "Toolbar", "Scrollbar"],
                application_level_json["shapes"],
            )
        )

        container_level_json = labeled_json.copy()
        container_level_json["shapes"] = list(
            filter(
                lambda x: x["label"]
                in [
                    "TabActive",
                    "TabInactive",
                    "Sidebar",
                    "Navbar",
                    "Container",
                    "Image",
                    "BrowserURLInput",
                ],
                container_level_json["shapes"],
            )
        )

        element_level_json = labeled_json.copy()
        element_level_json["shapes"] = list(
            filter(
                lambda x: x["label"]
                in [
                    "WebIcon",
                    "Icon",
                    "Switch",
                    "BtnSq",
                    "BtnPill",
                    "BtnCirc",
                    "CheckboxChecked",
                    "CheckboxUnchecked",
                    "RadiobtnSelected",
                    "RadiobtnUnselected",
                    "TextInput",
                    "Dropdown",
                    "Link",
                ],
                element_level_json["shapes"],
            )
        )

        text_level_json = labeled_json.copy()
        text_level_json["shapes"] = list(filter(lambda x: x["label"] == "Text", text_level_json["shapes"]))

        json.dump(
            top_level_json,
            open(f"{prefix}_TopLevel/{file}", "w"),
            indent=2,
        )

        json.dump(
            application_level_json,
            open(f"{prefix}_ApplicationLevel/{file}", "w"),
            indent=2,
        )

        json.dump(
            container_level_json,
            open(f"{prefix}_ContainerLevel/{file}", "w"),
            indent=2,
        )

        json.dump(
            element_level_json,
            open(f"{prefix}_ElementLevel/{file}", "w"),
            indent=2,
        )

        json.dump(
            text_level_json,
            open(f"{prefix}_TextLevel/{file}", "w"),
            indent=2,
        )

    else:
        copyfile(f"{prefix}/{file}", f"{prefix}_TopLevel/{file}")
        copyfile(f"{prefix}/{file}", f"{prefix}_ApplicationLevel/{file}")
        copyfile(f"{prefix}/{file}", f"{prefix}_ContainerLevel/{file}")
        copyfile(f"{prefix}/{file}", f"{prefix}_ElementLevel/{file}")
        copyfile(f"{prefix}/{file}", f"{prefix}_TextLevel/{file}")