Skip to content

Commit c42d151

Browse files
committed
WIP to objects file
1 parent 2d894fa commit c42d151

File tree

2 files changed

+159
-19
lines changed

2 files changed

+159
-19
lines changed

labelbox/data/serialization/ndjson/label.py

Lines changed: 54 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
from pydantic import BaseModel
77

8-
from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation, VideoClassificationAnnotation
8+
from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation, VideoClassificationAnnotation, VideoObjectAnnotation
99
from ...annotation_types.collection import LabelCollection, LabelGenerator
1010
from ...annotation_types.data import ImageData, TextData, VideoData
1111
from ...annotation_types.label import Label
@@ -33,6 +33,7 @@ def to_common(self) -> LabelGenerator:
3333
def from_common(cls,
3434
data: LabelCollection) -> Generator["NDLabel", None, None]:
3535
for label in data:
36+
# print(f"label in data:{label}")
3637
yield from cls._create_non_video_annotations(label)
3738
yield from cls._create_video_annotations(label)
3839

@@ -45,8 +46,10 @@ def _generate_annotations(
4546
for data_row_id, annotations in grouped_annotations.items():
4647
annots = []
4748
for annotation in annotations:
49+
4850
if isinstance(annotation, NDObjectType.__args__):
4951
annots.append(NDObject.to_common(annotation))
52+
#TODO: have a check on if the return type needs to extend or not
5053
elif isinstance(annotation, NDClassificationType.__args__):
5154
annots.extend(NDClassification.to_common(annotation))
5255
elif isinstance(annotation,
@@ -65,7 +68,7 @@ def _infer_media_type(
6568
types = {type(annotation) for annotation in annotations}
6669
if TextEntity in types:
6770
return TextData
68-
elif VideoClassificationAnnotation in types:
71+
elif VideoClassificationAnnotation in types or VideoObjectAnnotation in types:
6972
return VideoData
7073
else:
7174
return ImageData
@@ -83,26 +86,67 @@ def _get_consecutive_frames(
8386
def _create_video_annotations(
8487
cls, label: Label
8588
) -> Generator[Union[NDChecklistSubclass, NDRadioSubclass], None, None]:
89+
8690
video_annotations = defaultdict(list)
8791
for annot in label.annotations:
88-
if isinstance(annot, VideoClassificationAnnotation):
92+
if isinstance(
93+
annot,
94+
(VideoClassificationAnnotation, VideoObjectAnnotation)):
8995
video_annotations[annot.feature_schema_id].append(annot)
9096

97+
#break this into two groups, classifications, and then objects
9198
for annotation_group in video_annotations.values():
9299
consecutive_frames = cls._get_consecutive_frames(
93100
sorted([annotation.frame for annotation in annotation_group]))
94-
annotation = annotation_group[0]
95-
frames_data = []
96-
for frames in consecutive_frames:
97-
frames_data.append({'start': frames[0], 'end': frames[-1]})
98-
annotation.extra.update({'frames': frames_data})
99-
yield NDClassification.from_common(annotation, label.data)
101+
102+
if isinstance(annotation_group[0], VideoClassificationAnnotation):
103+
annotation = annotation_group[0]
104+
frames_data = []
105+
for frames in consecutive_frames:
106+
frames_data.append({'start': frames[0], 'end': frames[-1]})
107+
annotation.extra.update({'frames': frames_data})
108+
yield NDClassification.from_common(annotation, label.data)
109+
110+
elif isinstance(annotation_group[0], VideoObjectAnnotation):
111+
segments = []
112+
for start_frame, end_frame in consecutive_frames:
113+
segment = []
114+
for annotation in annotation_group:
115+
if annotation.keyframe and start_frame <= annotation.frame <= end_frame:
116+
segment.append(annotation)
117+
segments.append(segment)
118+
119+
print(segments[0], "\n")
120+
print(segments[1], "\n")
121+
print(consecutive_frames)
122+
yield NDObject.from_common(segments, label.data)
123+
# segments = []
124+
# seg_frames = []
125+
# for cframes in consecutive_frames:
126+
# seg_frames.append(cframes[0])
127+
# seg_frames.append(cframes[1])
128+
# print(seg_frames)
129+
# for annotation in annotation_group:
130+
# if annotation.frame in seg_frames:
131+
# segments.append(annotation)
132+
# # if annotation.keyframe:
133+
# # segments.append(annotation)
134+
# # print(consecutive_frames)
135+
# #TODO: current issue is that the way the code is written doesn't account for
136+
# #which frames are consecutive. maybe we should just have list of segments
137+
# annotations = []
138+
# for annotation in segments:
139+
# annotations.append(
140+
# NDObject.from_common(annotation, label.data))
141+
# yield annotations[0]
142+
# yield {}
100143

101144
@classmethod
102145
def _create_non_video_annotations(cls, label: Label):
103146
non_video_annotations = [
104147
annot for annot in label.annotations
105-
if not isinstance(annot, VideoClassificationAnnotation)
148+
if not isinstance(annot, (VideoClassificationAnnotation,
149+
VideoObjectAnnotation))
106150
]
107151
for annotation in non_video_annotations:
108152
if isinstance(annotation, ClassificationAnnotation):

labelbox/data/serialization/ndjson/objects.py

Lines changed: 105 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,18 @@
33
from typing import Any, Dict, List, Tuple, Union
44
import base64
55
import numpy as np
6+
from py import process
67

78
from pydantic import BaseModel
89
from PIL import Image
910

11+
from labelbox.data.annotation_types.data.video import VideoData
12+
1013
from ...annotation_types.data import ImageData, TextData, MaskData
1114
from ...annotation_types.ner import TextEntity
1215
from ...annotation_types.types import Cuid
1316
from ...annotation_types.geometry import Rectangle, Polygon, Line, Point, Mask
14-
from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation
17+
from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation, VideoObjectAnnotation
1518
from .classification import NDSubclassification, NDSubclassificationType
1619
from .base import DataRow, NDAnnotation
1720

@@ -118,6 +121,72 @@ def from_common(cls, rectangle: Rectangle,
118121
classifications=classifications)
119122

120123

124+
class NDFrameRectangle(BaseModel):
125+
frame: int
126+
bbox: Bbox
127+
128+
129+
class NDSegment(BaseModel):
130+
keyframes: List[NDFrameRectangle]
131+
132+
133+
class NDSegments(NDBaseObject):
134+
segments: List[NDSegment]
135+
136+
# uuid: str
137+
# schema_id: str
138+
139+
# def process_segment(self, segment: List[VideoObjectAnnotation]):
140+
# """
141+
# We only care about the annotation.value and frame once we make it here
142+
# """
143+
# # for annotation in segment:
144+
# return [{
145+
# "frame":
146+
# annotation.frame,
147+
# "bbox":
148+
# Bbox(top=annotation.value.start.y,
149+
# left=annotation.value.start.x,
150+
# height=annotation.value.end.y - annotation.value.start.y,
151+
# width=annotation.value.end.x - annotation.value.start.x)
152+
# } for annotation in segment]
153+
154+
def to_common(self):
155+
pass
156+
157+
@classmethod
158+
def from_common(cls, segments: List[VideoObjectAnnotation], data: VideoData,
159+
feature_schema_id: Cuid, extra: Dict[str,
160+
Any]) -> "NDSegments":
161+
print(f"\nWE MADE IT HERE TO SEGMENTS\n")
162+
# for segment in segments:
163+
# print("\nSEGMENT\n", segment)
164+
# processed_segment = cls.process_segment(cls, segment)
165+
166+
# segments = [{
167+
# "keyframes": [{
168+
# "frame": 3,
169+
# "bbox": Bbox(top=0, left=0, height=1, width=1)
170+
# }]
171+
# }, {
172+
# "keyframes": [{
173+
# "frame": 5,
174+
# "bbox": Bbox(top=0, left=0, height=3, width=5)
175+
# }]
176+
# }]
177+
178+
segments = [{"keyframes": segment} for segment in segments]
179+
180+
print("wew\n", segments[0], "sss\n")
181+
182+
a = cls(segments=segments,
183+
dataRow=DataRow(id=data.uid),
184+
schema_id=feature_schema_id,
185+
uuid=extra.get('uuid'))
186+
print("A\n", a, "\nZ")
187+
return a
188+
189+
121190
class _URIMask(BaseModel):
122191
instanceURI: str
123192
colorRGB: Tuple[int, int, int]
@@ -211,6 +280,21 @@ def from_common(
211280
cls, annotation: ObjectAnnotation, data: Union[ImageData, TextData]
212281
) -> Union[NDLine, NDPoint, NDPolygon, NDRectangle, NDMask, NDTextEntity]:
213282
obj = cls.lookup_object(annotation)
283+
284+
#if it is video segments
285+
if (obj == NDSegments):
286+
print("hello i am ndsegment")
287+
#look into segment of segments
288+
#look into annotations of segment
289+
#check and validate that there are no subclasses
290+
#new method for processing segments ?
291+
292+
return obj.from_common(
293+
annotation,
294+
data,
295+
feature_schema_id=annotation[0][0].feature_schema_id,
296+
extra=annotation[0][0].extra)
297+
214298
subclasses = [
215299
NDSubclassification.from_common(annot)
216300
for annot in annotation.classifications
@@ -221,20 +305,32 @@ def from_common(
221305

222306
@staticmethod
223307
def lookup_object(annotation: ObjectAnnotation) -> "NDObjectType":
224-
result = {
225-
Line: NDLine,
226-
Point: NDPoint,
227-
Polygon: NDPolygon,
228-
Rectangle: NDRectangle,
229-
Mask: NDMask,
230-
TextEntity: NDTextEntity
231-
}.get(type(annotation.value))
308+
if isinstance(annotation, list):
309+
result = NDSegments
310+
else:
311+
result = {
312+
Line: NDLine,
313+
Point: NDPoint,
314+
Polygon: NDPolygon,
315+
Rectangle: NDRectangle,
316+
Mask: NDMask,
317+
TextEntity: NDTextEntity
318+
}.get(type(annotation.value))
232319
if result is None:
233320
raise TypeError(
234321
f"Unable to convert object to MAL format. `{type(annotation.value)}`"
235322
)
236323
return result
237324

325+
# @staticmethod
326+
# def process_segments(segments: List[List[VideoObjectAnnotation]]):
327+
# for segment in segments:
328+
# for annotation in segment:
329+
# subclasses = [
330+
# NDSubclassification.from_common(annot)
331+
# for annot in annotation.classifications
332+
# ]
333+
238334

239335
NDObjectType = Union[NDLine, NDPolygon, NDPoint, NDRectangle, NDMask,
240336
NDTextEntity]

0 commit comments

Comments
 (0)