Skip to content

Commit ca1c6f9

Browse files
LeoEATLindenLeoEATLinden
authored
Merge pull request #5 from superannotateai/restructuring
Restructuring
2 parents a0f8911 + 041d68d commit ca1c6f9

20 files changed

+643
-87
lines changed

‎src/superannotate_databricks_connector/schemas/comment.py‎

Lines changed: 49 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@
55
FloatType,
66
BooleanType,
77
MapType,
8-
ArrayType
8+
ArrayType,
9+
IntegerType
910
)
1011

12+
from .shapes import get_bbox_schema
1113

12-
def get_comment_schema():
13-
comment_schema = StructType([
14+
15+
def get_vector_comment_schema():
16+
return StructType([
1417
StructField("correspondence",
1518
ArrayType(MapType(
1619
StringType(),
@@ -23,12 +26,53 @@ def get_comment_schema():
2326
StructField("createdBy", MapType(
2427
StringType(),
2528
StringType()),
29+
True),
30+
StructField("creationType", StringType(), True),
31+
StructField("updatedAt", StringType(), True),
32+
StructField("updatedBy", MapType(
33+
StringType(),
34+
StringType()),
35+
True)
36+
])
37+
38+
39+
def get_video_timestamp_schema():
40+
return StructType([
41+
StructField("timestamp", IntegerType(), True),
42+
StructField("points", get_bbox_schema(), True)
43+
])
44+
45+
46+
def get_video_comment_parameter_schema():
47+
return StructType([
48+
StructField("start", IntegerType(), True),
49+
StructField("end", IntegerType, True),
50+
StructField("timestamps", ArrayType(
51+
get_video_timestamp_schema()), True)
52+
])
53+
54+
55+
def get_video_comment_schema():
56+
return StructType([
57+
StructField("correspondence",
58+
ArrayType(MapType(
59+
StringType(),
60+
StringType())),
2661
True),
62+
StructField("start", IntegerType(), True),
63+
StructField("end", IntegerType(), True),
64+
StructField("createdAt", StringType(), True),
65+
StructField("createdBy", MapType(
66+
StringType(),
67+
StringType()),
68+
True),
2769
StructField("creationType", StringType(), True),
2870
StructField("updatedAt", StringType(), True),
2971
StructField("updatedBy", MapType(
3072
StringType(),
3173
StringType()),
32-
True)
74+
True),
75+
StructField("parameters",
76+
ArrayType(get_video_comment_parameter_schema()), True)
77+
3378
])
34-
return comment_schema

‎src/superannotate_databricks_connector/schemas/shapes.py‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
from pyspark.sql.types import (
2+
StructType,
3+
StructField,
4+
FloatType,
5+
ArrayType
6+
)
7+
8+
9+
def get_bbox_schema():
10+
"""
11+
Defines the schema of a bounding box
12+
13+
Args:
14+
None
15+
16+
Returns:
17+
StructType: Schema of a bbox
18+
"""
19+
return StructType([
20+
StructField("x1", FloatType(), True),
21+
StructField("y1", FloatType(), True),
22+
StructField("x2", FloatType(), True),
23+
StructField("y2", FloatType(), True)
24+
])
25+
26+
27+
def get_rbbox_schema():
28+
"""
29+
Defines the schema of a rotated bounding box
30+
this contains one point for each corned
31+
32+
Args:
33+
None
34+
35+
Returns:
36+
StructType: Schema of a bbox
37+
"""
38+
return StructType([
39+
StructField("x1", FloatType(), True),
40+
StructField("y1", FloatType(), True),
41+
StructField("x2", FloatType(), True),
42+
StructField("y2", FloatType(), True),
43+
StructField("x3", FloatType(), True),
44+
StructField("y3", FloatType(), True),
45+
StructField("x4", FloatType(), True),
46+
StructField("y5", FloatType(), True)
47+
])
48+
49+
50+
def get_point_schema():
51+
"""
52+
Defines the schema of a point
53+
54+
Args:
55+
None
56+
57+
Returns:
58+
StructType: Schema of a point
59+
"""
60+
return StructType([
61+
StructField("x", FloatType(), True),
62+
StructField("y", FloatType(), True)
63+
])
64+
65+
66+
def get_cuboid_schema():
67+
"""
68+
Defines the schema of a cuboid (3d bounding box)
69+
70+
Args:
71+
None
72+
73+
Returns:
74+
StructType: Schema of a cuboid
75+
"""
76+
return StructType([
77+
StructField("f1", get_point_schema(), True),
78+
StructField("f2", get_point_schema(), True),
79+
StructField("r1", get_point_schema(), True),
80+
StructField("r2", get_point_schema(), True)
81+
])
82+
83+
84+
def get_ellipse_schema():
85+
"""
86+
Defines the schema of an ellipse
87+
88+
Args:
89+
None
90+
91+
Returns:
92+
StructType: Schema of an ellipse
93+
"""
94+
return StructType([
95+
StructField("cx", FloatType(), True),
96+
StructField("cy", FloatType(), True),
97+
StructField("rx", FloatType(), True),
98+
StructField("ty", FloatType(), True),
99+
StructField("angle", FloatType(), True)
100+
])
101+
102+
103+
def get_polygon_schema():
104+
"""
105+
Defines the schema of a polygon. It contains a shell as well
106+
as excluded points
107+
108+
Args:
109+
None
110+
111+
Returns:
112+
StructType: Schema of a polygon with holes
113+
"""
114+
return StructType([
115+
StructField("points", ArrayType(FloatType()), True),
116+
StructField("exclude", ArrayType(ArrayType(FloatType())), True)
117+
])
118+
119+
120+
def get_polyline_schema():
121+
"""
122+
Defines the schema of a polyline
123+
A simple array of float
124+
125+
Args:
126+
None
127+
128+
Returns:
129+
ArrayType: Schema of a polygon with holes
130+
"""
131+
return ArrayType(FloatType())

‎src/superannotate_databricks_connector/schemas/tag.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
from pyspark.sql.types import (
2+
StructType,
3+
StructField,
4+
StringType,
5+
IntegerType,
6+
MapType,
7+
ArrayType,
8+
)
9+
10+
11+
def get_tag_schema():
12+
schema = StructType([
13+
StructField("instance_type", StringType(), True),
14+
StructField("classId", IntegerType(), True),
15+
StructField("probability", IntegerType(), True),
16+
StructField("attributes", ArrayType(MapType(StringType(),
17+
StringType())),
18+
True),
19+
StructField("createdAt", StringType(), True),
20+
StructField("createdBy", MapType(StringType(), StringType()), True),
21+
StructField("creationType", StringType(), True),
22+
StructField("updatedAt", StringType(), True),
23+
StructField("updatedBy", MapType(StringType(), StringType()), True),
24+
StructField("className", StringType(), True)])
25+
return schema

‎src/superannotate_databricks_connector/schemas/vector_schema.py‎

Lines changed: 21 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -3,45 +3,35 @@
33
StructField,
44
StringType,
55
IntegerType,
6-
FloatType,
76
BooleanType,
87
MapType,
98
ArrayType
109
)
11-
from .comment import get_comment_schema
12-
13-
14-
def get_point_schema():
15-
point_schema = StructType([
16-
StructField("x", FloatType(), True),
17-
StructField("y", FloatType(), True)
18-
])
19-
return point_schema
20-
21-
22-
def get_cuboid_schema():
23-
cuboid_points_schema = StructType([
24-
StructField("f1", get_point_schema(), True),
25-
StructField("f2", get_point_schema(), True),
26-
StructField("r1", get_point_schema(), True),
27-
StructField("r2", get_point_schema(), True)
28-
])
29-
return cuboid_points_schema
10+
from .comment import get_vector_comment_schema
11+
from .shapes import (
12+
get_point_schema,
13+
get_cuboid_schema,
14+
get_bbox_schema,
15+
get_ellipse_schema,
16+
get_polygon_schema,
17+
get_polyline_schema,
18+
get_rbbox_schema
19+
)
20+
from .tag import get_tag_schema
3021

3122

3223
def get_vector_instance_schema():
33-
instance_schema = StructType([
24+
return StructType([
3425
StructField("instance_type", StringType(), True),
3526
StructField("classId", IntegerType(), True),
3627
StructField("probability", IntegerType(), True),
37-
StructField("bbox_points", MapType(StringType(), FloatType()), True),
38-
StructField("polygon_points", ArrayType(FloatType()), True),
39-
StructField("polygon_exclude", ArrayType(ArrayType(FloatType())),
40-
True),
41-
StructField("cuboid_points", get_cuboid_schema(), True),
42-
StructField("ellipse_points", MapType(StringType(), FloatType()),
43-
True),
44-
StructField("point_points", MapType(StringType(), FloatType()), True),
28+
StructField("bbox", get_bbox_schema(), True),
29+
StructField("rbbox", get_rbbox_schema(), True),
30+
StructField("polygon", get_polygon_schema()),
31+
StructField("cuboid", get_cuboid_schema(), True),
32+
StructField("ellipse", get_ellipse_schema(), True),
33+
StructField("polyline", get_polyline_schema(), True),
34+
StructField("point", get_point_schema(), True),
4535
StructField("groupId", IntegerType(), True),
4636
StructField("locked", BooleanType(), True),
4737
StructField("attributes", ArrayType(MapType(StringType(),
@@ -56,24 +46,6 @@ def get_vector_instance_schema():
5646
StructField("updatedBy", MapType(StringType(), StringType()), True),
5747
StructField("className", StringType(), True)
5848
])
59-
return instance_schema
60-
61-
62-
def get_vector_tag_schema():
63-
schema = StructType([
64-
StructField("instance_type", StringType(), True),
65-
StructField("classId", IntegerType(), True),
66-
StructField("probability", IntegerType(), True),
67-
StructField("attributes", ArrayType(MapType(StringType(),
68-
StringType())),
69-
True),
70-
StructField("createdAt", StringType(), True),
71-
StructField("createdBy", MapType(StringType(), StringType()), True),
72-
StructField("creationType", StringType(), True),
73-
StructField("updatedAt", StringType(), True),
74-
StructField("updatedBy", MapType(StringType(), StringType()), True),
75-
StructField("className", StringType(), True)])
76-
return schema
7749

7850

7951
def get_vector_schema():
@@ -90,7 +62,7 @@ def get_vector_schema():
9062
StructField("instances", ArrayType(get_vector_instance_schema()),
9163
True),
9264
StructField("bounding_boxes", ArrayType(IntegerType()), True),
93-
StructField("comments", ArrayType(get_comment_schema()), True),
94-
StructField("tags", ArrayType(get_vector_tag_schema()), True)
65+
StructField("comments", ArrayType(get_vector_comment_schema()), True),
66+
StructField("tags", ArrayType(get_tag_schema()), True)
9567
])
9668
return schema

0 commit comments

Comments
 (0)