Add stubs for pycocotools (#9086)

Co-authored-by: Alex Waygood <[email protected]>

Author: hoel-bagard

stubs/pycocotools/METADATA.toml

@@ -0,0 +1,4 @@
+version = "2.0.*"
+
+[tool.stubtest]
+ignore_missing_stub = false

stubs/pycocotools/pycocotools/__init__.pyi

@@ -0,0 +1,6 @@
+from typing_extensions import TypedDict
+
+# Unused in this module, but imported in multiple submodules.
+class _EncodedRLE(TypedDict):  # noqa: Y049
+    size: list[int]
+    counts: str | bytes

stubs/pycocotools/pycocotools/coco.pyi

@@ -0,0 +1,94 @@
+from _typeshed import Incomplete
+from collections.abc import Collection, Sequence
+from pathlib import Path
+from typing import Generic, TypeVar, overload
+from typing_extensions import Literal, TypeAlias, TypedDict
+
+from . import _EncodedRLE
+
+# TODO: Use numpy types when #5768 is resolved.
+# import numpy as np
+# import numpy.typing as npt
+
+PYTHON_VERSION: Incomplete
+_NDArray: TypeAlias = Incomplete
+
+class _Image(TypedDict):
+    id: int
+    width: int
+    height: int
+    file_name: str
+
+_TPolygonSegmentation: TypeAlias = list[list[float]]
+
+class _RLE(TypedDict):
+    size: list[int]
+    counts: list[int]
+
+class _Annotation(TypedDict):
+    id: int
+    image_id: int
+    category_id: int
+    segmentation: _TPolygonSegmentation | _RLE | _EncodedRLE
+    area: float
+    bbox: list[float]
+    iscrowd: int
+
+_TSeg = TypeVar("_TSeg", _TPolygonSegmentation, _RLE, _EncodedRLE)
+
+class _AnnotationG(TypedDict, Generic[_TSeg]):
+    id: int
+    image_id: int
+    category_id: int
+    segmentation: _TSeg
+    area: float
+    bbox: list[float]
+    iscrowd: int
+
+class _Category(TypedDict):
+    id: int
+    name: str
+    supercategory: str
+
+class _Dataset(TypedDict):
+    images: list[_Image]
+    annotations: list[_Annotation]
+    categories: list[_Category]
+
+class COCO:
+    anns: dict[int, _Annotation]
+    dataset: _Dataset
+    cats: dict[int, _Category]
+    imgs: dict[int, _Image]
+    imgToAnns: dict[int, list[_Annotation]]
+    catToImgs: dict[int, list[int]]
+    def __init__(self, annotation_file: str | Path | None = ...) -> None: ...
+    def createIndex(self) -> None: ...
+    def info(self) -> None: ...
+    def getAnnIds(
+        self,
+        imgIds: Collection[int] | int = ...,
+        catIds: Collection[int] | int = ...,
+        areaRng: Sequence[float] = ...,
+        iscrowd: bool | None = ...,
+    ) -> list[int]: ...
+    def getCatIds(
+        self, catNms: Collection[str] | str = ..., supNms: Collection[str] | str = ..., catIds: Collection[int] | int = ...
+    ) -> list[int]: ...
+    def getImgIds(self, imgIds: Collection[int] | int = ..., catIds: list[int] | int = ...) -> list[int]: ...
+    def loadAnns(self, ids: Collection[int] | int = ...) -> list[_Annotation]: ...
+    def loadCats(self, ids: Collection[int] | int = ...) -> list[_Category]: ...
+    def loadImgs(self, ids: Collection[int] | int = ...) -> list[_Image]: ...
+    def showAnns(self, anns: Sequence[_Annotation], draw_bbox: bool = ...) -> None: ...
+    def loadRes(self, resFile: str) -> COCO: ...
+    def download(self, tarDir: str | None = ..., imgIds: Collection[int] = ...) -> Literal[-1] | None: ...
+    def loadNumpyAnnotations(self, data: _NDArray) -> list[_Annotation]: ...
+    # def loadNumpyAnnotations(self, data: npt.NDArray[np.float64]) -> list[_Annotation]: ...
+    @overload
+    def annToRLE(self, ann: _AnnotationG[_RLE]) -> _RLE: ...
+    @overload
+    def annToRLE(self, ann: _AnnotationG[_EncodedRLE]) -> _EncodedRLE: ...
+    @overload
+    def annToRLE(self, ann: _AnnotationG[_TPolygonSegmentation]) -> _EncodedRLE: ...
+    def annToMask(self, ann: _Annotation) -> _NDArray: ...
+    # def annToMask(self, ann: _Annotation) -> npt.NDArray[np.uint8]: ...

stubs/pycocotools/pycocotools/cocoeval.pyi

@@ -0,0 +1,65 @@
+from _typeshed import Incomplete
+from typing_extensions import Literal, TypeAlias, TypedDict
+
+from .coco import COCO
+
+# TODO: Use numpy types when #5768 is resolved.
+# import numpy as np
+# import numpy.typing as npt
+
+_NDArray: TypeAlias = Incomplete
+_TIOU: TypeAlias = Literal["segm", "bbox", "keypoints"]
+
+class _EvaluationResult(TypedDict):
+    image_id: int
+    category_id: int
+    aRng: list[int]
+    maxDet: int
+    dtIds: list[int]
+    gtIds: list[int]
+    dtMatches: _NDArray
+    # dtMatches: npt.NDArray[np.float64]
+    gtMatches: _NDArray
+    # gtMatches: npt.NDArray[np.float64]
+    dtScores: list[float]
+    gtIgnore: _NDArray
+    # gtIgnore: npt.NDArray[np.float64]
+    dtIgnore: _NDArray
+    # dtIgnore: npt.NDArray[np.float64]
+
+class COCOeval:
+    cocoGt: COCO
+    cocoDt: COCO
+    evalImgs: list[_EvaluationResult]
+    eval: _EvaluationResult
+    params: Params
+    stats: _NDArray
+    # stats: npt.NDArray[np.float64]
+    ious: dict[tuple[int, int], list[float]]
+    def __init__(self, cocoGt: COCO | None = ..., cocoDt: COCO | None = ..., iouType: _TIOU = ...) -> None: ...
+    def evaluate(self) -> None: ...
+    def computeIoU(self, imgId: int, catId: int) -> list[float]: ...
+    def computeOks(self, imgId: int, catId: int) -> _NDArray: ...
+    # def computeOks(self, imgId: int, catId: int) -> npt.NDArray[np.float64]: ...
+    def evaluateImg(self, imgId: int, catId: int, aRng: list[int], maxDet: int) -> _EvaluationResult: ...
+    def accumulate(self, p: Params | None = ...) -> None: ...
+    def summarize(self) -> None: ...
+
+class Params:
+    imgIds: list[int]
+    catIds: list[int]
+    iouThrs: _NDArray
+    # iouThrs: npt.NDArray[np.float64]
+    recThrs: _NDArray
+    # recThrs: npt.NDArray[np.float64]
+    maxDets: list[int]
+    areaRng: list[int]
+    areaRngLbl: list[str]
+    useCats: int
+    kpt_oks_sigmas: _NDArray
+    # kpt_oks_sigmas: npt.NDArray[np.float64]
+    iouType: _TIOU
+    useSegm: int | None
+    def __init__(self, iouType: _TIOU = ...) -> None: ...
+    def setDetParams(self) -> None: ...
+    def setKpParams(self) -> None: ...

stubs/pycocotools/pycocotools/mask.pyi

@@ -0,0 +1,31 @@
+from _typeshed import Incomplete
+from typing import Any, overload
+from typing_extensions import TypeAlias
+
+from . import _EncodedRLE
+
+# TODO: Use numpy types when #5768 is resolved.
+# import numpy as np
+# import numpy.typing as npt
+
+_NPUInt32: TypeAlias = Incomplete  # np.uint32
+_NDArrayUInt8: TypeAlias = Incomplete  # npt.NDArray[np.uint8]
+_NDArrayUInt32: TypeAlias = Incomplete  # npt.NDArray[np.uint32]
+_NDArrayFloat64: TypeAlias = Incomplete  # npt.NDArray[np.float64]
+
+def iou(
+    dt: _NDArrayUInt32 | list[float] | list[_EncodedRLE],
+    gt: _NDArrayUInt32 | list[float] | list[_EncodedRLE],
+    pyiscrowd: list[int] | _NDArrayUInt8,
+) -> list[Any] | _NDArrayFloat64: ...
+def merge(rleObjs: list[_EncodedRLE], intersect: int = ...) -> _EncodedRLE: ...
+
+# ignore an "overlapping overloads" error due to _NDArrayInt32 being an alias for `Incomplete` for now
+@overload
+def frPyObjects(pyobj: _NDArrayUInt32 | list[list[int]] | list[_EncodedRLE], h: int, w: int) -> list[_EncodedRLE]: ...  # type: ignore[misc]
+@overload
+def frPyObjects(pyobj: list[int] | _EncodedRLE, h: int, w: int) -> _EncodedRLE: ...
+def encode(bimask: _NDArrayUInt8) -> _EncodedRLE: ...
+def decode(rleObjs: _EncodedRLE) -> _NDArrayUInt8: ...
+def area(rleObjs: _EncodedRLE) -> _NPUInt32: ...
+def toBbox(rleObjs: _EncodedRLE) -> _NDArrayFloat64: ...