diff --git a/rbproto/python/shot_pb2.py b/rbproto/python/shot_pb2.py index 86866c9..9dd02cf 100644 --- a/rbproto/python/shot_pb2.py +++ b/rbproto/python/shot_pb2.py @@ -13,7 +13,7 @@ _sym_db = _symbol_database.Default() -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\nshot.proto\x12\x0fserialized_shot\"?\n\x03\x42ox\x12\x0c\n\x04left\x18\x01 \x01(\x02\x12\x0b\n\x03top\x18\x02 \x01(\x02\x12\r\n\x05width\x18\x03 \x01(\x02\x12\x0e\n\x06height\x18\x04 \x01(\x02\"\x1d\n\x05Point\x12\t\n\x01x\x18\x01 \x01(\x02\x12\t\n\x01y\x18\x02 \x01(\x02\"\x7f\n\rBallDetection\x12.\n\x0eplane_position\x18\x01 \x01(\x0b\x32\x16.serialized_shot.Point\x12(\n\nannotation\x18\x02 \x01(\x0b\x32\x14.serialized_shot.Box\x12\x14\n\x0cinterpolated\x18\x03 \x01(\x08\"T\n\x10RLEBallDetection\x12\x31\n\tdetection\x18\x01 \x01(\x0b\x32\x1e.serialized_shot.BallDetection\x12\r\n\x05\x63ount\x18\x02 \x01(\r\"L\n\x13RLEDetectionHistory\x12\x35\n\ndetections\x18\x01 \x03(\x0b\x32!.serialized_shot.RLEBallDetection\"F\n\x10\x44\x65tectionHistory\x12\x32\n\ndetections\x18\x01 \x03(\x0b\x32\x1e.serialized_shot.BallDetection\"\xfc\x01\n\rCollisionInfo\x12\x0e\n\x06source\x18\x01 \x01(\r\x12M\n\x10\x62\x61ll_identifiers\x18\x02 \x03(\x0b\x32\x33.serialized_shot.CollisionInfo.BallIdentifiersEntry\x12\x17\n\x0fwall_identifier\x18\x03 \x01(\r\x12\x13\n\x0b\x66rame_index\x18\x04 \x01(\r\x12\x0e\n\x06static\x18\x05 \x01(\x08\x1aN\n\x14\x42\x61llIdentifiersEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.serialized_shot.Point:\x02\x38\x01\"\xcc\x02\n\x04Path\x12\x13\n\x0bstart_frame\x18\x01 \x01(\r\x12\x11\n\tend_frame\x18\x02 \x01(\r\x12\x37\n\ndetections\x18\x03 \x01(\x0b\x32!.serialized_shot.DetectionHistoryH\x00\x12>\n\x0erle_detections\x18\x04 \x01(\x0b\x32$.serialized_shot.RLEDetectionHistoryH\x00\x12\x15\n\x0bnot_present\x18\x05 \x01(\x08H\x00\x12\x11\n\tis_static\x18\x06 \x01(\x08\x12\x32\n\nstart_info\x18\x07 \x01(\x0b\x32\x1e.serialized_shot.CollisionInfo\x12\x30\n\x08\x65nd_info\x18\x08 \x01(\x0b\x32\x1e.serialized_shot.CollisionInfoB\x13\n\x11\x64\x65tection_history\"R\n\x11IdentifierHistory\x12\x17\n\x0f\x62\x61ll_identifier\x18\x01 \x01(\r\x12$\n\x05paths\x18\x02 \x03(\x0b\x32\x15.serialized_shot.Path\"H\n\x04Shot\x12@\n\x14identifier_histories\x18\x03 \x03(\x0b\x32\".serialized_shot.IdentifierHistoryb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\nshot.proto\x12\x0fserialized_shot\"?\n\x03\x42ox\x12\x0c\n\x04left\x18\x01 \x01(\x02\x12\x0b\n\x03top\x18\x02 \x01(\x02\x12\r\n\x05width\x18\x03 \x01(\x02\x12\x0e\n\x06height\x18\x04 \x01(\x02\"\x1d\n\x05Point\x12\t\n\x01x\x18\x01 \x01(\x02\x12\t\n\x01y\x18\x02 \x01(\x02\"\x7f\n\rBallDetection\x12.\n\x0eplane_position\x18\x01 \x01(\x0b\x32\x16.serialized_shot.Point\x12(\n\nannotation\x18\x02 \x01(\x0b\x32\x14.serialized_shot.Box\x12\x14\n\x0cinterpolated\x18\x03 \x01(\x08\"T\n\x10RLEBallDetection\x12\x31\n\tdetection\x18\x01 \x01(\x0b\x32\x1e.serialized_shot.BallDetection\x12\r\n\x05\x63ount\x18\x02 \x01(\r\"L\n\x13RLEDetectionHistory\x12\x35\n\ndetections\x18\x01 \x03(\x0b\x32!.serialized_shot.RLEBallDetection\"F\n\x10\x44\x65tectionHistory\x12\x32\n\ndetections\x18\x01 \x03(\x0b\x32\x1e.serialized_shot.BallDetection\"\xfc\x01\n\rCollisionInfo\x12\x0e\n\x06source\x18\x01 \x01(\r\x12M\n\x10\x62\x61ll_identifiers\x18\x02 \x03(\x0b\x32\x33.serialized_shot.CollisionInfo.BallIdentifiersEntry\x12\x17\n\x0fwall_identifier\x18\x03 \x01(\r\x12\x13\n\x0b\x66rame_index\x18\x04 \x01(\r\x12\x0e\n\x06static\x18\x05 \x01(\x08\x1aN\n\x14\x42\x61llIdentifiersEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.serialized_shot.Point:\x02\x38\x01\"\xcc\x02\n\x04Path\x12\x13\n\x0bstart_frame\x18\x01 \x01(\r\x12\x11\n\tend_frame\x18\x02 \x01(\r\x12\x37\n\ndetections\x18\x03 \x01(\x0b\x32!.serialized_shot.DetectionHistoryH\x00\x12>\n\x0erle_detections\x18\x04 \x01(\x0b\x32$.serialized_shot.RLEDetectionHistoryH\x00\x12\x15\n\x0bnot_present\x18\x05 \x01(\x08H\x00\x12\x11\n\tis_static\x18\x06 \x01(\x08\x12\x32\n\nstart_info\x18\x07 \x01(\x0b\x32\x1e.serialized_shot.CollisionInfo\x12\x30\n\x08\x65nd_info\x18\x08 \x01(\x0b\x32\x1e.serialized_shot.CollisionInfoB\x13\n\x11\x64\x65tection_history\"R\n\x11IdentifierHistory\x12\x17\n\x0f\x62\x61ll_identifier\x18\x01 \x01(\r\x12$\n\x05paths\x18\x02 \x03(\x0b\x32\x15.serialized_shot.Path\"j\n\x12KeyBallIdentifiers\x12\x10\n\x08\x63ue_ball\x18\x01 \x01(\r\x12\x13\n\x0bobject_ball\x18\x02 \x01(\r\x12\x13\n\x0btarget_ball\x18\x03 \x01(\r\x12\x18\n\x10\x63ontact_sequence\x18\x04 \x03(\r\"\x80\x01\n\x04Shot\x12@\n\x14identifier_histories\x18\x03 \x03(\x0b\x32\".serialized_shot.IdentifierHistory\x12\x36\n\tkey_balls\x18\x04 \x01(\x0b\x32#.serialized_shot.KeyBallIdentifiersb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -42,6 +42,8 @@ if _descriptor._USE_C_DESCRIPTORS == False: _globals['_PATH']._serialized_end=1080 _globals['_IDENTIFIERHISTORY']._serialized_start=1082 _globals['_IDENTIFIERHISTORY']._serialized_end=1164 - _globals['_SHOT']._serialized_start=1166 - _globals['_SHOT']._serialized_end=1238 + _globals['_KEYBALLIDENTIFIERS']._serialized_start=1166 + _globals['_KEYBALLIDENTIFIERS']._serialized_end=1272 + _globals['_SHOT']._serialized_start=1275 + _globals['_SHOT']._serialized_end=1403 # @@protoc_insertion_point(module_scope) diff --git a/rbproto/python/shot_pb2.pyi b/rbproto/python/shot_pb2.pyi index 2fbf17a..495a5de 100644 --- a/rbproto/python/shot_pb2.pyi +++ b/rbproto/python/shot_pb2.pyi @@ -104,8 +104,22 @@ class IdentifierHistory(_message.Message): paths: _containers.RepeatedCompositeFieldContainer[Path] def __init__(self, ball_identifier: _Optional[int] = ..., paths: _Optional[_Iterable[_Union[Path, _Mapping]]] = ...) -> None: ... +class KeyBallIdentifiers(_message.Message): + __slots__ = ["cue_ball", "object_ball", "target_ball", "contact_sequence"] + CUE_BALL_FIELD_NUMBER: _ClassVar[int] + OBJECT_BALL_FIELD_NUMBER: _ClassVar[int] + TARGET_BALL_FIELD_NUMBER: _ClassVar[int] + CONTACT_SEQUENCE_FIELD_NUMBER: _ClassVar[int] + cue_ball: int + object_ball: int + target_ball: int + contact_sequence: _containers.RepeatedScalarFieldContainer[int] + def __init__(self, cue_ball: _Optional[int] = ..., object_ball: _Optional[int] = ..., target_ball: _Optional[int] = ..., contact_sequence: _Optional[_Iterable[int]] = ...) -> None: ... + class Shot(_message.Message): - __slots__ = ["identifier_histories"] + __slots__ = ["identifier_histories", "key_balls"] IDENTIFIER_HISTORIES_FIELD_NUMBER: _ClassVar[int] + KEY_BALLS_FIELD_NUMBER: _ClassVar[int] identifier_histories: _containers.RepeatedCompositeFieldContainer[IdentifierHistory] - def __init__(self, identifier_histories: _Optional[_Iterable[_Union[IdentifierHistory, _Mapping]]] = ...) -> None: ... + key_balls: KeyBallIdentifiers + def __init__(self, identifier_histories: _Optional[_Iterable[_Union[IdentifierHistory, _Mapping]]] = ..., key_balls: _Optional[_Union[KeyBallIdentifiers, _Mapping]] = ...) -> None: ... diff --git a/rbproto/shot.proto b/rbproto/shot.proto index 83be5ff..08106aa 100644 --- a/rbproto/shot.proto +++ b/rbproto/shot.proto @@ -3,10 +3,10 @@ syntax = "proto3"; package serialized_shot; message Box { - float left = 1; - float top = 2; - float width = 3; - float height = 4; + float left = 1; + float top = 2; + float width = 3; + float height = 4; } message Point { @@ -15,14 +15,14 @@ message Point { } message BallDetection { - Point plane_position = 1; - Box annotation = 2; - bool interpolated = 3; + Point plane_position = 1; + Box annotation = 2; + bool interpolated = 3; } message RLEBallDetection { - BallDetection detection = 1; - uint32 count = 2; + BallDetection detection = 1; + uint32 count = 2; } message RLEDetectionHistory { @@ -58,7 +58,18 @@ message IdentifierHistory { uint32 ball_identifier = 1; repeated Path paths = 2; } +message KeyBallIdentifiers { + uint32 cue_ball = 1; + uint32 object_ball = 2; + uint32 target_ball = 3; + + // For now this will just be cue->object/target->target + // Long term this will potentially represent a linked list + // of all balls in a shot. + repeated uint32 contact_sequence = 4; +} message Shot { repeated IdentifierHistory identifier_histories = 3; + KeyBallIdentifiers key_balls = 4; } diff --git a/rbproto/ts/shot.ts b/rbproto/ts/shot.ts index 6c04053..e3399a9 100644 --- a/rbproto/ts/shot.ts +++ b/rbproto/ts/shot.ts @@ -815,8 +815,121 @@ function _decodeIdentifierHistory(bb: ByteBuffer): IdentifierHistory { return message; } +export interface KeyBallIdentifiers { + cue_ball?: number; + object_ball?: number; + target_ball?: number; + contact_sequence?: number[]; +} + +export function encodeKeyBallIdentifiers( + message: KeyBallIdentifiers, +): Uint8Array { + let bb = popByteBuffer(); + _encodeKeyBallIdentifiers(message, bb); + return toUint8Array(bb); +} + +function _encodeKeyBallIdentifiers( + message: KeyBallIdentifiers, + bb: ByteBuffer, +): void { + // optional uint32 cue_ball = 1; + let $cue_ball = message.cue_ball; + if ($cue_ball !== undefined) { + writeVarint32(bb, 8); + writeVarint32(bb, $cue_ball); + } + + // optional uint32 object_ball = 2; + let $object_ball = message.object_ball; + if ($object_ball !== undefined) { + writeVarint32(bb, 16); + writeVarint32(bb, $object_ball); + } + + // optional uint32 target_ball = 3; + let $target_ball = message.target_ball; + if ($target_ball !== undefined) { + writeVarint32(bb, 24); + writeVarint32(bb, $target_ball); + } + + // repeated uint32 contact_sequence = 4; + let array$contact_sequence = message.contact_sequence; + if (array$contact_sequence !== undefined) { + let packed = popByteBuffer(); + for (let value of array$contact_sequence) { + writeVarint32(packed, value); + } + writeVarint32(bb, 34); + writeVarint32(bb, packed.offset); + writeByteBuffer(bb, packed); + pushByteBuffer(packed); + } +} + +export function decodeKeyBallIdentifiers( + binary: Uint8Array, +): KeyBallIdentifiers { + return _decodeKeyBallIdentifiers(wrapByteBuffer(binary)); +} + +function _decodeKeyBallIdentifiers(bb: ByteBuffer): KeyBallIdentifiers { + let message: KeyBallIdentifiers = {} as any; + + end_of_message: while (!isAtEnd(bb)) { + let tag = readVarint32(bb); + + switch (tag >>> 3) { + case 0: + break end_of_message; + + // optional uint32 cue_ball = 1; + case 1: { + message.cue_ball = readVarint32(bb) >>> 0; + break; + } + + // optional uint32 object_ball = 2; + case 2: { + message.object_ball = readVarint32(bb) >>> 0; + break; + } + + // optional uint32 target_ball = 3; + case 3: { + message.target_ball = readVarint32(bb) >>> 0; + break; + } + + // repeated uint32 contact_sequence = 4; + case 4: { + let values = + message.contact_sequence || (message.contact_sequence = []); + if ((tag & 7) === 2) { + let outerLimit = pushTemporaryLength(bb); + while (!isAtEnd(bb)) { + values.push(readVarint32(bb) >>> 0); + } + bb.limit = outerLimit; + } else { + values.push(readVarint32(bb) >>> 0); + } + break; + } + + default: + skipUnknownField(bb, tag & 7); + } + } + + return message; +} + export interface Shot { identifier_histories?: IdentifierHistory[]; + key_balls?: KeyBallIdentifiers; } export function encodeShot(message: Shot): Uint8Array { @@ -838,6 +951,17 @@ function _encodeShot(message: Shot, bb: ByteBuffer): void { pushByteBuffer(nested); } } + + // optional KeyBallIdentifiers key_balls = 4; + let $key_balls = message.key_balls; + if ($key_balls !== undefined) { + writeVarint32(bb, 34); + let nested = popByteBuffer(); + _encodeKeyBallIdentifiers($key_balls, nested); + writeVarint32(bb, nested.limit); + writeByteBuffer(bb, nested); + pushByteBuffer(nested); + } } export function decodeShot(binary: Uint8Array): Shot { @@ -864,6 +988,14 @@ function _decodeShot(bb: ByteBuffer): Shot { break; } + // optional KeyBallIdentifiers key_balls = 4; + case 4: { + let limit = pushTemporaryLength(bb); + message.key_balls = _decodeKeyBallIdentifiers(bb); + bb.limit = limit; + break; + } + default: skipUnknownField(bb, tag & 7); }