<commit_before>from typing import Optional

from torch.nn.modules.batchnorm import BatchNorm2d
from torch.nn.modules.instancenorm import InstanceNorm2d
from torchvision.models.optical_flow import RAFT
from torchvision.models.optical_flow.raft import _raft, BottleneckBlock, ResidualBlock
from torchvision.prototype.transforms import RaftEval
from torchvision.transforms.functional import InterpolationMode

from .._api import WeightsEnum
from .._api import Weights
from .._utils import handle_legacy_interface


__all__ = (
    "RAFT",
    "raft_large",
    "raft_small",
    "Raft_Large_Weights",
    "Raft_Small_Weights",
)


_COMMON_META = {"interpolation": InterpolationMode.BILINEAR}


class Raft_Large_Weights(WeightsEnum):
    C_T_V1 = Weights(
        # Chairs + Things, ported from original paper repo (raft-things.pth)
        url="https://download.pytorch.org/models/raft_large_C_T_V1-22a6c225.pth",
        transforms=RaftEval,
        meta={
            **_COMMON_META,
            "recipe": "https://github.com/princeton-vl/RAFT",
            "sintel_train_cleanpass_epe": 1.4411,
            "sintel_train_finalpass_epe": 2.7894,
            "kitti_train_per_image_epe": 5.0172,
            "kitti_train_f1-all": 17.4506,
        },
    )

    C_T_V2 = Weights(
        # Chairs + Things
        url="https://download.pytorch.org/models/raft_large_C_T_V2-1bb1363a.pth",
        transforms=RaftEval,
        meta={
            **_COMMON_META,
            "recipe": "https://github.com/pytorch/vision/tree/main/references/optical_flow",
            "sintel_train_cleanpass_epe": 1.3822,
            "sintel_train_finalpass_epe": 2.7161,
            "kitti_train_per_image_epe": 4.5118,
            "kitti_train_f1-all": 16.0679,
        },
    )

    # C_T_SKHT_V1 = Weights(
    #     # Chairs + Things + Sintel fine-tuning, i.e.:
    #     # Chairs + Things + (Sintel + Kitti + HD1K + Things_clean)
    #     # Corresponds to the C+T+S+K+H on paper with fine-tuning on Sintel
    #     url="",
    #     transforms=RaftEval,
    #     meta={
    #         "recipe": "",
    #         "epe": -1234,
    #     },
    # )

    # C_T_SKHT_K_V1 = Weights(
    #     # Chairs + Things + Sintel fine-tuning + Kitti fine-tuning i.e.:
    #     # Chairs + Things + (Sintel + Kitti + HD1K + Things_clean) + Kitti
    #     # Same as CT_SKHT with extra fine-tuning on Kitti
    #     # Corresponds to the C+T+S+K+H on paper with fine-tuning on Sintel and then on Kitti
    #     url="",
    #     transforms=RaftEval,
    #     meta={
    #         "recipe": "",
    #         "epe": -1234,
    #     },<commit_msg>Pre-trained weights for raft-large, fine-tuned on sintel (#5079)<commit_after>from typing import Optional

from torch.nn.modules.batchnorm import BatchNorm2d
from torch.nn.modules.instancenorm import InstanceNorm2d
from torchvision.models.optical_flow import RAFT
from torchvision.models.optical_flow.raft import _raft, BottleneckBlock, ResidualBlock
from torchvision.prototype.transforms import RaftEval
from torchvision.transforms.functional import InterpolationMode

from .._api import WeightsEnum
from .._api import Weights
from .._utils import handle_legacy_interface


__all__ = (
    "RAFT",
    "raft_large",
    "raft_small",
    "Raft_Large_Weights",
    "Raft_Small_Weights",
)


_COMMON_META = {"interpolation": InterpolationMode.BILINEAR}


class Raft_Large_Weights(WeightsEnum):
    C_T_V1 = Weights(
        # Chairs + Things, ported from original paper repo (raft-things.pth)
        url="https://download.pytorch.org/models/raft_large_C_T_V1-22a6c225.pth",
        transforms=RaftEval,
        meta={
            **_COMMON_META,
            "recipe": "https://github.com/princeton-vl/RAFT",
            "sintel_train_cleanpass_epe": 1.4411,
            "sintel_train_finalpass_epe": 2.7894,
            "kitti_train_per_image_epe": 5.0172,
            "kitti_train_f1-all": 17.4506,
        },
    )

    C_T_V2 = Weights(
        # Chairs + Things
        url="https://download.pytorch.org/models/raft_large_C_T_V2-1bb1363a.pth",
        transforms=RaftEval,
        meta={
            **_COMMON_META,
            "recipe": "https://github.com/pytorch/vision/tree/main/references/optical_flow",
            "sintel_train_cleanpass_epe": 1.3822,
            "sintel_train_finalpass_epe": 2.7161,
            "kitti_train_per_image_epe": 4.5118,
            "kitti_train_f1-all": 16.0679,
        },
    )

    C_T_SKHT_V1 = Weights(
        # Chairs + Things + Sintel fine-tuning, ported from original paper repo (raft-sintel.pth)
        url="https://download.pytorch.org/models/raft_large_C_T_SKHT_V1-0b8c9e55.pth",
        transforms=RaftEval,
        meta={
            **_COMMON_META,
            "recipe": "https://github.com/princeton-vl/RAFT",
            "sintel_test_cleanpass_epe": 1.94,
            "sintel_test_finalpass_epe": 3.18,
        },
    )

    C_T_SKHT_V2 = Weights(
        # Chairs + Things + Sintel fine-tuning, i.e.:
        # Chairs + Things + (Sintel + Kitti + HD1K + Things_clean)
        # Corresponds to the C+T+S+K+H on paper with fine-tuning on Sintel
        url="https://download.pytorch.org/models/raft_large_C_T_SKHT_V2-ff5fadd5.pth",
        transforms=RaftEval,
        meta={
            **_COMMON_META,
            "recipe": "https://github.com/pytorch/vision/tree/main/references/optical_flow",
            "sintel_test_cleanpass_epe": 1.819,
            "sintel_test_finalpass_epe": 3.067,
        },
    )

    # C_T_SKHT_K_V1 = Weights(
    #     # Chairs + Things + Sintel fine-tuning + Kitti fine-tuning i.e.:
    #     # Chairs + Things + (Sintel + Kitti + HD1K + Things_clean) + Kitti
    #     # Same as CT_SKHT with extra fine-tuning on Kitti
    #     # Corresponds to the C+T+S+K+H on paper with fine-tuning on Sintel and then on Kitti
    #     url="",
    #     transforms=RaftEval,
    #     meta={
    #         "recipe": "",
    #         "epe": -1234,
    #     },
{'fraction_non_alphanumeric': 0.1144103338366046, 'fraction_numerical': 0.02851870491528267, 'mean_word_length': 2.4845119812974867, 'pattern_counts': {'":': 34, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 12, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '972698', 'n_tokens_mistral': 2254, 'n_tokens_neox': 2140, 'n_words': 369}