jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def table_struct_recog(image, THRESHOLD_PROBA):
    '''
    Table structure recognition using DEtect-object TRansformer pre-trained on 1 million tables
    '''

    feature_extractor = DetrFeatureExtractor(do_resize=True, size=1000, max_size=1000)
    encoding = feature_extractor(image, return_tensors="pt")

    model = DetrForObjectDetection.from_pretrained("SalML/DETR-table-structure-recognition")
    with torch.no_grad():
        outputs = model(**encoding)

    probas = outputs.logits.softmax(-1)[0, :, :-1]
    keep = probas.max(-1).values > THRESHOLD_PROBA

    target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)
    postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
    bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]

    return (model, probas[keep], bboxes_scaled)
//  [python]
def table_struct_recog(image, THRESHOLD_PROBA):
    '''
    Table structure recognition using DEtect-object TRansformer pre-trained on 1 million tables
    '''

    feature_extractor = DetrImageProcessor(do_resize=True, size=1000, max_size=1000)
    encoding = feature_extractor(image, return_tensors="pt")

    model = DetrForObjectDetection.from_pretrained("SalML/DETR-table-structure-recognition")
    with torch.no_grad():
        outputs = model(**encoding)

    probas = outputs.logits.softmax(-1)[0, :, :-1]
    keep = probas.max(-1).values > THRESHOLD_PROBA

    target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)
    postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
    bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]

    return (model, probas[keep], bboxes_scaled)
//  [python]
def runTATR(image):
        feature_extractor = DetrFeatureExtractor()
        encoding = feature_extractor(image, return_tensors="pt")
        encoding.keys()

        model = TableTransformerForObjectDetection.from_pretrained(
            "microsoft/table-transformer-structure-recognition")

        with torch.no_grad():
            outputs = model(**encoding)
        target_sizes = [image.size[::-1]]

        results = feature_extractor.post_process_object_detection(outputs, threshold=0.1, target_sizes=target_sizes)[0]

        return results, model, target_sizes
//  [python]
def runTATR(image):
        feature_extractor = DetrImageProcessor()
        encoding = feature_extractor(image, return_tensors="pt")
        encoding.keys()

        model = TableTransformerForObjectDetection.from_pretrained(
            "microsoft/table-transformer-structure-recognition")

        with torch.no_grad():
            outputs = model(**encoding)
        target_sizes = [image.size[::-1]]

        results = feature_extractor.post_process_object_detection(outputs, threshold=0.1, target_sizes=target_sizes)[0]

        return results, model, target_sizes
//  [python]
def initialize(
        self,
        model: Union[str, Path, TableTransformerForObjectDetection] = None,
        device: Optional[str] = "cuda" if torch.cuda.is_available() else "cpu",
    ):
        """Loads the donut model using the specified parameters"""
        self.device = device
        self.feature_extractor = DetrFeatureExtractor()

        try:
            logger.info("Loading the table structure model ...")
            self.model = TableTransformerForObjectDetection.from_pretrained(model)
            self.model.eval()

        except EnvironmentError:
            logger.critical("Failed to initialize the model.")
            logger.critical("Ensure that the model is correct")
            raise ImportError(
                "Review the parameters to initialize a UnstructuredTableTransformerModel obj",
            )
        self.model.to(device)
//  [python]
def initialize(
        self,
        model: Union[str, Path, TableTransformerForObjectDetection] = None,
        device: Optional[str] = "cuda" if torch.cuda.is_available() else "cpu",
    ):
        """Loads the donut model using the specified parameters"""
        self.device = device
        self.feature_extractor = DetrImageProcessor()

        try:
            logger.info("Loading the table structure model ...")
            self.model = TableTransformerForObjectDetection.from_pretrained(model)
            self.model.eval()

        except EnvironmentError:
            logger.critical("Failed to initialize the model.")
            logger.critical("Ensure that the model is correct")
            raise ImportError(
                "Review the parameters to initialize a UnstructuredTableTransformerModel obj",
            )
        self.model.to(device)
//  [python]
def __init__(self) -> None:
        self.model = DetrFeatureExtractor()
//  [python]
def __init__(self) -> None:
        self.model = DetrImageProcessor()
//  [python]
def convert_glpn_checkpoint(checkpoint_path, pytorch_dump_folder_path, push_to_hub=False, model_name=None):
    """
    Copy/paste/tweak model's weights to our GLPN structure.
    """

    # load GLPN configuration (Segformer-B4 size)
    config = GLPNConfig(hidden_sizes=[64, 128, 320, 512], decoder_hidden_size=64, depths=[3, 8, 27, 3])

    # load feature extractor (only resize + rescale)
    feature_extractor = GLPNFeatureExtractor()

    # prepare image
    image = prepare_img()
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values

    logger.info("Converting model...")

    # load original state dict
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))

    # rename keys
    state_dict = rename_keys(state_dict)

    # key and value matrices need special treatment
    read_in_k_v(state_dict, config)

    # create HuggingFace model and load state dict
    model = GLPNForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    # forward pass
    outputs = model(pixel_values)
    predicted_depth = outputs.predicted_depth

    # verify output
    if model_name is not None:
        if "nyu" in model_name:
            expected_slice = torch.tensor(
                [[4.4147, 4.0873, 4.0673], [3.7890, 3.2881, 3.1525], [3.7674, 3.5423, 3.4913]]
            )
        elif "kitti" in model_name:
            expected_slice = torch.tensor(
                [[3.4291, 2.7865, 2.5151], [3.2841, 2.7021, 2.3502], [3.1147, 2.4625, 2.2481]]
            )
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        expected_shape = torch.Size([1, 480, 640])

        assert predicted_depth.shape == expected_shape
        assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4)
        print("Looks ok!")

    # finally, push to hub if required
    if push_to_hub:
        logger.info("Pushing model and feature extractor to the hub...")
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )
        feature_extractor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add feature extractor",
            use_temp_dir=True,
        )
//  [python]
def convert_glpn_checkpoint(checkpoint_path, pytorch_dump_folder_path, push_to_hub=False, model_name=None):
    """
    Copy/paste/tweak model's weights to our GLPN structure.
    """

    # load GLPN configuration (Segformer-B4 size)
    config = GLPNConfig(hidden_sizes=[64, 128, 320, 512], decoder_hidden_size=64, depths=[3, 8, 27, 3])

    # load feature extractor (only resize + rescale)
    feature_extractor = GLPNImageProcessor()

    # prepare image
    image = prepare_img()
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values

    logger.info("Converting model...")

    # load original state dict
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))

    # rename keys
    state_dict = rename_keys(state_dict)

    # key and value matrices need special treatment
    read_in_k_v(state_dict, config)

    # create HuggingFace model and load state dict
    model = GLPNForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    # forward pass
    outputs = model(pixel_values)
    predicted_depth = outputs.predicted_depth

    # verify output
    if model_name is not None:
        if "nyu" in model_name:
            expected_slice = torch.tensor(
                [[4.4147, 4.0873, 4.0673], [3.7890, 3.2881, 3.1525], [3.7674, 3.5423, 3.4913]]
            )
        elif "kitti" in model_name:
            expected_slice = torch.tensor(
                [[3.4291, 2.7865, 2.5151], [3.2841, 2.7021, 2.3502], [3.1147, 2.4625, 2.2481]]
            )
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        expected_shape = torch.Size([1, 480, 640])

        assert predicted_depth.shape == expected_shape
        assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4)
        print("Looks ok!")

    # finally, push to hub if required
    if push_to_hub:
        logger.info("Pushing model and feature extractor to the hub...")
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )
        feature_extractor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add feature extractor",
            use_temp_dir=True,
        )
//  [python]
def convert_glpn_checkpoint(
    checkpoint_path, pytorch_dump_folder_path, push_to_hub=False, model_name=None
):
    """
    Copy/paste/tweak model's weights to our GLPN structure.
    """

    # load GLPN configuration (Segformer-B4 size)
    config = GLPNConfig(
        hidden_sizes=[64, 128, 320, 512], decoder_hidden_size=64, depths=[3, 8, 27, 3]
    )

    # load feature extractor (only resize + rescale)
    feature_extractor = GLPNFeatureExtractor()

    # prepare image
    image = prepare_img()
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values

    logger.info("Converting model...")

    # load original state dict
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))

    # rename keys
    state_dict = rename_keys(state_dict)

    # key and value matrices need special treatment
    read_in_k_v(state_dict, config)

    # create HuggingFace model and load state dict
    model = GLPNForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    # forward pass
    outputs = model(pixel_values)
    predicted_depth = outputs.predicted_depth

    # verify output
    if model_name is not None:
        if "nyu" in model_name:
            expected_slice = torch.tensor(
                [
                    [4.4147, 4.0873, 4.0673],
                    [3.7890, 3.2881, 3.1525],
                    [3.7674, 3.5423, 3.4913],
                ]
            )
        elif "kitti" in model_name:
            expected_slice = torch.tensor(
                [
                    [3.4291, 2.7865, 2.5151],
                    [3.2841, 2.7021, 2.3502],
                    [3.1147, 2.4625, 2.2481],
                ]
            )
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        expected_shape = torch.Size([1, 480, 640])

        assert predicted_depth.shape == expected_shape
        assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4)
        print("Looks ok!")

    # finally, push to hub if required
    if push_to_hub:
        logger.info("Pushing model and feature extractor to the hub...")
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )
        feature_extractor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add feature extractor",
            use_temp_dir=True,
        )
//  [python]
def convert_glpn_checkpoint(
    checkpoint_path, pytorch_dump_folder_path, push_to_hub=False, model_name=None
):
    """
    Copy/paste/tweak model's weights to our GLPN structure.
    """

    # load GLPN configuration (Segformer-B4 size)
    config = GLPNConfig(
        hidden_sizes=[64, 128, 320, 512], decoder_hidden_size=64, depths=[3, 8, 27, 3]
    )

    # load feature extractor (only resize + rescale)
    feature_extractor = GLPNImageProcessor()

    # prepare image
    image = prepare_img()
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values

    logger.info("Converting model...")

    # load original state dict
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))

    # rename keys
    state_dict = rename_keys(state_dict)

    # key and value matrices need special treatment
    read_in_k_v(state_dict, config)

    # create HuggingFace model and load state dict
    model = GLPNForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    # forward pass
    outputs = model(pixel_values)
    predicted_depth = outputs.predicted_depth

    # verify output
    if model_name is not None:
        if "nyu" in model_name:
            expected_slice = torch.tensor(
                [
                    [4.4147, 4.0873, 4.0673],
                    [3.7890, 3.2881, 3.1525],
                    [3.7674, 3.5423, 3.4913],
                ]
            )
        elif "kitti" in model_name:
            expected_slice = torch.tensor(
                [
                    [3.4291, 2.7865, 2.5151],
                    [3.2841, 2.7021, 2.3502],
                    [3.1147, 2.4625, 2.2481],
                ]
            )
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        expected_shape = torch.Size([1, 480, 640])

        assert predicted_depth.shape == expected_shape
        assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4)
        print("Looks ok!")

    # finally, push to hub if required
    if push_to_hub:
        logger.info("Pushing model and feature extractor to the hub...")
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )
        feature_extractor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add feature extractor",
            use_temp_dir=True,
        )
//  [python]
def convert_glpn_checkpoint(checkpoint_path, pytorch_dump_folder_path, push_to_hub=False, model_name=None):
    """
    Copy/paste/tweak model's weights to our GLPN structure.
    """

    # load GLPN configuration (Segformer-B4 size)
    config = GLPNConfig(hidden_sizes=[64, 128, 320, 512], decoder_hidden_size=64, depths=[3, 8, 27, 3])

    # load image processor (only resize + rescale)
    image_processor = GLPNFeatureExtractor()

    # prepare image
    image = prepare_img()
    pixel_values = image_processor(images=image, return_tensors="pt").pixel_values

    logger.info("Converting model...")

    # load original state dict
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))

    # rename keys
    state_dict = rename_keys(state_dict)

    # key and value matrices need special treatment
    read_in_k_v(state_dict, config)

    # create HuggingFace model and load state dict
    model = GLPNForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    # forward pass
    outputs = model(pixel_values)
    predicted_depth = outputs.predicted_depth

    # verify output
    if model_name is not None:
        if "nyu" in model_name:
            expected_slice = torch.tensor(
                [[4.4147, 4.0873, 4.0673], [3.7890, 3.2881, 3.1525], [3.7674, 3.5423, 3.4913]]
            )
        elif "kitti" in model_name:
            expected_slice = torch.tensor(
                [[3.4291, 2.7865, 2.5151], [3.2841, 2.7021, 2.3502], [3.1147, 2.4625, 2.2481]]
            )
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        expected_shape = torch.Size([1, 480, 640])

        assert predicted_depth.shape == expected_shape
        assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4)
        print("Looks ok!")

    # finally, push to hub if required
    if push_to_hub:
        logger.info("Pushing model and image processor to the hub...")
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )
        image_processor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add image processor",
            use_temp_dir=True,
        )
//  [python]
def convert_glpn_checkpoint(checkpoint_path, pytorch_dump_folder_path, push_to_hub=False, model_name=None):
    """
    Copy/paste/tweak model's weights to our GLPN structure.
    """

    # load GLPN configuration (Segformer-B4 size)
    config = GLPNConfig(hidden_sizes=[64, 128, 320, 512], decoder_hidden_size=64, depths=[3, 8, 27, 3])

    # load image processor (only resize + rescale)
    image_processor = GLPNImageProcessor()

    # prepare image
    image = prepare_img()
    pixel_values = image_processor(images=image, return_tensors="pt").pixel_values

    logger.info("Converting model...")

    # load original state dict
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))

    # rename keys
    state_dict = rename_keys(state_dict)

    # key and value matrices need special treatment
    read_in_k_v(state_dict, config)

    # create HuggingFace model and load state dict
    model = GLPNForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    # forward pass
    outputs = model(pixel_values)
    predicted_depth = outputs.predicted_depth

    # verify output
    if model_name is not None:
        if "nyu" in model_name:
            expected_slice = torch.tensor(
                [[4.4147, 4.0873, 4.0673], [3.7890, 3.2881, 3.1525], [3.7674, 3.5423, 3.4913]]
            )
        elif "kitti" in model_name:
            expected_slice = torch.tensor(
                [[3.4291, 2.7865, 2.5151], [3.2841, 2.7021, 2.3502], [3.1147, 2.4625, 2.2481]]
            )
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        expected_shape = torch.Size([1, 480, 640])

        assert predicted_depth.shape == expected_shape
        assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4)
        print("Looks ok!")

    # finally, push to hub if required
    if push_to_hub:
        logger.info("Pushing model and image processor to the hub...")
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )
        image_processor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add image processor",
            use_temp_dir=True,
        )
//  [python]
def check_training(self, config, input_values, *args):
        model = TFHubertForCTC(config)

        # freeze feature encoder
        model.freeze_feature_extractor()

        input_values = input_values[:3]

        input_lengths = tf.constant([input_values.shape[-1] // i for i in [4, 2, 1]])
        max_length_labels = model.hubert._get_feat_extract_output_lengths(input_lengths)
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        length_mask = tf.sequence_mask(input_lengths, dtype=tf.float32)

        input_values = input_values * length_mask

        pad_size = max(max_length_labels) - labels.shape[1]
        labels = tf.pad(labels, ((0, 0), (0, pad_size)), constant_values=-100)

        loss = model(input_values, labels=labels, training=True).loss

        self.parent.assertFalse(tf.math.is_inf(loss))
//  [python]
def check_training(self, config, input_values, *args):
        model = TFHubertForCTC(config)

        # freeze feature encoder
        model.freeze_feature_encoder()

        input_values = input_values[:3]

        input_lengths = tf.constant([input_values.shape[-1] // i for i in [4, 2, 1]])
        max_length_labels = model.hubert._get_feat_extract_output_lengths(input_lengths)
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        length_mask = tf.sequence_mask(input_lengths, dtype=tf.float32)

        input_values = input_values * length_mask

        pad_size = max(max_length_labels) - labels.shape[1]
        labels = tf.pad(labels, ((0, 0), (0, pad_size)), constant_values=-100)

        loss = model(input_values, labels=labels, training=True).loss

        self.parent.assertFalse(tf.math.is_inf(loss))
//  [python]
def check_training(self, config, input_values, *args):
        model = TFHubertForCTC(config)

        # freeze feature extractor
        model.freeze_feature_extractor()

        input_values = input_values[:3]

        input_lengths = tf.constant([input_values.shape[-1] // i for i in [4, 2, 1]])
        max_length_labels = model.hubert._get_feat_extract_output_lengths(input_lengths)
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        length_mask = tf.sequence_mask(input_lengths, dtype=tf.float32)

        input_values = input_values * length_mask

        pad_size = max(max_length_labels) - labels.shape[1]
        labels = tf.pad(labels, ((0, 0), (0, pad_size)), constant_values=-100)

        loss = model(input_values, labels=labels, training=True).loss

        self.parent.assertFalse(tf.math.is_inf(loss))
//  [python]
def test_small_model_pt(self):
        small_model = "hf-internal-testing/tiny-random-VideoMAEForVideoClassification"
        small_feature_extractor = VideoMAEFeatureExtractor(
            size={"shortest_edge": 10}, crop_size={"height": 10, "width": 10}
        )
        video_classifier = pipeline(
            "video-classification", model=small_model, feature_extractor=small_feature_extractor, frame_sampling_rate=4
        )

        video_file_path = hf_hub_download(repo_id="nateraw/video-demo", filename="archery.mp4", repo_type="dataset")
        outputs = video_classifier(video_file_path, top_k=2)
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}],
        )

        outputs = video_classifier(
            [
                video_file_path,
                video_file_path,
            ],
            top_k=2,
        )
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
                [{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}],
                [{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}],
            ],
        )
//  [python]
def test_small_model_pt(self):
        small_model = "hf-internal-testing/tiny-random-VideoMAEForVideoClassification"
        small_feature_extractor = VideoMAEImageProcessor(
            size={"shortest_edge": 10}, crop_size={"height": 10, "width": 10}
        )
        video_classifier = pipeline(
            "video-classification", model=small_model, feature_extractor=small_feature_extractor, frame_sampling_rate=4
        )

        video_file_path = hf_hub_download(repo_id="nateraw/video-demo", filename="archery.mp4", repo_type="dataset")
        outputs = video_classifier(video_file_path, top_k=2)
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}],
        )

        outputs = video_classifier(
            [
                video_file_path,
                video_file_path,
            ],
            top_k=2,
        )
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
                [{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}],
                [{"score": 0.5199, "label": "LABEL_0"}, {"score": 0.4801, "label": "LABEL_1"}],
            ],
        )
//  [python]
def default_feature_extractor(self):
        # logits were tested with a different mean and std, so we use the same here
        return (
            VideoMAEFeatureExtractor(image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
            if is_vision_available()
            else None
        )
//  [python]
def default_feature_extractor(self):
        # logits were tested with a different mean and std, so we use the same here
        return (
            VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
            if is_vision_available()
            else None
        )
//  [python]
def generate_video_lmdb(pretrain_path, video_path, frame_no):

	configuration = VideoMAEConfig()
	feature_extractor = VideoMAEFeatureExtractor(configuration)
	feature_extractor = feature_extractor.from_pretrained(pretrain_path)

	frames_folder = 'video_frames/' # a folder that restores frames of videos
	# {video1-1.jpg, video1-2.jpg, video1-3.jpg, video1-4.jpg, video1-5.jpg} refer to 5 frames extracted from video1
	lmdb_path = 'video_frames.lmdb' # to input

	all_video = [x.replace('.mp4', '.jpg') for x in os.listdir(video_path)]
	video_num = len(all_video)
	isdir = os.path.isdir(lmdb_path)
	lmdb_env = lmdb.open(lmdb_path, subdir=isdir, map_size=video_num * np.zeros((3, 224, 224)).nbytes*10,
		readonly=False, meminit=False, map_async=True)
	txn = lmdb_env.begin(write=True)
	write_frequency = 100
	for idx, video in enumerate(tqdm.tqdm(all_video)):
		temp = []
		video_id = int(video)
		for frame_idx in range(1, frame_no+1):
			img = np.array(transform(Image.open(frames_folder+str(video_id)+'-'+str(frame_idx)+'.jpg').convert('RGB')))		
			temp.append(img)
		temp = list(np.array(temp))
		temp = feature_extractor(temp, return_tensors="pt").pixel_values[0]
		temp = LMDB_VIDEO(np.array(temp))
		txn.put(u'{}'.format(video_id).encode('ascii'), pickle.dumps(temp))
		if video_id % write_frequency == 0 and idx != 0:
			txn.commit()
			txn = lmdb_env.begin(write=True)
	txn.commit()
	keys = [u'{}'.format(k).encode('ascii') for k in range(video_num)]
	with lmdb_env.begin(write=True) as txn:
		txn.put(b'__keys__', pickle.dumps(keys))
		txn.put(b'__len__', pickle.dumps(len(keys)))
	print(len(keys))
	print("Flushing database ...")
	lmdb_env.sync()
	lmdb_env.close()
//  [python]
def generate_video_lmdb(pretrain_path, video_path, frame_no):

	configuration = VideoMAEConfig()
	feature_extractor = VideoMAEImageProcessor(configuration)
	feature_extractor = feature_extractor.from_pretrained(pretrain_path)

	frames_folder = 'video_frames/' # a folder that restores frames of videos
	# {video1-1.jpg, video1-2.jpg, video1-3.jpg, video1-4.jpg, video1-5.jpg} refer to 5 frames extracted from video1
	lmdb_path = 'video_frames.lmdb' # to input

	all_video = [x.replace('.mp4', '.jpg') for x in os.listdir(video_path)]
	video_num = len(all_video)
	isdir = os.path.isdir(lmdb_path)
	lmdb_env = lmdb.open(lmdb_path, subdir=isdir, map_size=video_num * np.zeros((3, 224, 224)).nbytes*10,
		readonly=False, meminit=False, map_async=True)
	txn = lmdb_env.begin(write=True)
	write_frequency = 100
	for idx, video in enumerate(tqdm.tqdm(all_video)):
		temp = []
		video_id = int(video)
		for frame_idx in range(1, frame_no+1):
			img = np.array(transform(Image.open(frames_folder+str(video_id)+'-'+str(frame_idx)+'.jpg').convert('RGB')))		
			temp.append(img)
		temp = list(np.array(temp))
		temp = feature_extractor(temp, return_tensors="pt").pixel_values[0]
		temp = LMDB_VIDEO(np.array(temp))
		txn.put(u'{}'.format(video_id).encode('ascii'), pickle.dumps(temp))
		if video_id % write_frequency == 0 and idx != 0:
			txn.commit()
			txn = lmdb_env.begin(write=True)
	txn.commit()
	keys = [u'{}'.format(k).encode('ascii') for k in range(video_num)]
	with lmdb_env.begin(write=True) as txn:
		txn.put(b'__keys__', pickle.dumps(keys))
		txn.put(b'__len__', pickle.dumps(len(keys)))
	print(len(keys))
	print("Flushing database ...")
	lmdb_env.sync()
	lmdb_env.close()
//  [python]
def default_image_processor(self):
        # logits were tested with a different mean and std, so we use the same here
        return (
            VideoMAEFeatureExtractor(image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
            if is_vision_available()
            else None
        )
//  [python]
def default_image_processor(self):
        # logits were tested with a different mean and std, so we use the same here
        return (
            VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
            if is_vision_available()
            else None
        )
//  [python]
def __init__(self, annotation, epoch_size=None, frame_transform=None, video_transform=None, clip_len=8, dataset=''):
        super(FrameSampleDataset2).__init__()

        #self.samples = get_samples(root)
        lines = []
        with open(annotation) as f:
            for line in f:
                split = line.strip().split(', ')
                if split[1] == 'label':
                    continue
                if dataset == 'SSm':
                    lines.append((split[0], str(int(split[1]))))
                else:
                    lines.append((split[0], str(int(split[1]) - 1)))

        print('start', lines[0], len(lines))
        self.samples = lines
        self.preprocess = VideoMAEFeatureExtractor(size={"height": 224, "width": 224}, do_center_crop=False)

        # Allow for temporal jittering
        if epoch_size is None:
            epoch_size = len(self.samples)
        self.epoch_size = epoch_size

        self.clip_len = clip_len
        self.frame_transform = self.preprocess#frame_transform
        self.video_transform = video_transform
//  [python]
def __init__(self, annotation, epoch_size=None, frame_transform=None, video_transform=None, clip_len=8, dataset=''):
        super(FrameSampleDataset2).__init__()

        #self.samples = get_samples(root)
        lines = []
        with open(annotation) as f:
            for line in f:
                split = line.strip().split(', ')
                if split[1] == 'label':
                    continue
                if dataset == 'SSm':
                    lines.append((split[0], str(int(split[1]))))
                else:
                    lines.append((split[0], str(int(split[1]) - 1)))

        print('start', lines[0], len(lines))
        self.samples = lines
        self.preprocess = VideoMAEImageProcessor(size= {"height": 224, "width": 224}, do_center_crop=False)

        # Allow for temporal jittering
        if epoch_size is None:
            epoch_size = len(self.samples)
        self.epoch_size = epoch_size

        self.clip_len = clip_len
        self.frame_transform = self.preprocess#frame_transform
        self.video_transform = video_transform
//  [python]
def check_ctc_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = SEWDForCTC(config=config)
        model.to(torch_device)
        model.train()

        # freeze feature encoder
        model.freeze_feature_extractor()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

            if max_length_labels[i] < labels.shape[-1]:
                # it's important that we make sure that target lenghts are at least
                # one shorter than logit lenghts to prevent -inf
                labels[i, max_length_labels[i] - 1 :] = -100

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()
//  [python]
def check_ctc_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = SEWDForCTC(config=config)
        model.to(torch_device)
        model.train()

        # freeze feature encoder
        model.freeze_feature_encoder()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

            if max_length_labels[i] < labels.shape[-1]:
                # it's important that we make sure that target lenghts are at least
                # one shorter than logit lenghts to prevent -inf
                labels[i, max_length_labels[i] - 1 :] = -100

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()
//  [python]
def check_ctc_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = SEWDForCTC(config=config)
        model.to(torch_device)
        model.train()

        # freeze feature encoder
        model.freeze_feature_extractor()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

            if max_length_labels[i] < labels.shape[-1]:
                # it's important that we make sure that target lengths are at least
                # one shorter than logit lengths to prevent -inf
                labels[i, max_length_labels[i] - 1 :] = -100

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()
//  [python]
def check_ctc_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = SEWDForCTC(config=config)
        model.to(torch_device)
        model.train()

        # freeze feature encoder
        model.freeze_feature_encoder()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

            if max_length_labels[i] < labels.shape[-1]:
                # it's important that we make sure that target lengths are at least
                # one shorter than logit lengths to prevent -inf
                labels[i, max_length_labels[i] - 1 :] = -100

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()
//  [python]
def test_integration_instance_segmentation(self):
        # load 2 images and corresponding annotations from the hub
        repo_id = "nielsr/image-segmentation-toy-data"
        image1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_image_1.png", repo_type="dataset")
        )
        image2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_image_2.png", repo_type="dataset")
        )
        annotation1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_annotation_1.png", repo_type="dataset")
        )
        annotation2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_annotation_2.png", repo_type="dataset")
        )

        # get instance segmentations and instance-to-segmentation mappings
        def get_instance_segmentation_and_mapping(annotation):
            instance_seg = np.array(annotation)[:, :, 1]
            class_id_map = np.array(annotation)[:, :, 0]
            class_labels = np.unique(class_id_map)

            # create mapping between instance IDs and semantic category IDs
            inst2class = {}
            for label in class_labels:
                instance_ids = np.unique(instance_seg[class_id_map == label])
                inst2class.update({i: label for i in instance_ids})

            return instance_seg, inst2class

        instance_seg1, inst2class1 = get_instance_segmentation_and_mapping(annotation1)
        instance_seg2, inst2class2 = get_instance_segmentation_and_mapping(annotation2)

        # create a feature extractor
        feature_extractor = MaskFormerFeatureExtractor(reduce_labels=True, ignore_index=255, size=(512, 512))

        # prepare the images and annotations
        inputs = feature_extractor(
            [image1, image2],
            [instance_seg1, instance_seg2],
            instance_id_to_semantic_id=[inst2class1, inst2class2],
            return_tensors="pt",
        )

        # verify the pixel values and pixel mask
        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))

        # verify the class labels
        self.assertEqual(len(inputs["class_labels"]), 2)
        self.assertTrue(torch.allclose(inputs["class_labels"][0], torch.tensor([30, 55])))
        self.assertTrue(torch.allclose(inputs["class_labels"][1], torch.tensor([4, 4, 23, 55])))

        # verify the mask labels
        self.assertEqual(len(inputs["mask_labels"]), 2)
        self.assertEqual(inputs["mask_labels"][0].shape, (2, 512, 512))
        self.assertEqual(inputs["mask_labels"][1].shape, (4, 512, 512))
        self.assertEquals(inputs["mask_labels"][0].sum().item(), 41527.0)
        self.assertEquals(inputs["mask_labels"][1].sum().item(), 26259.0)
//  [python]
def test_integration_instance_segmentation(self):
        # load 2 images and corresponding annotations from the hub
        repo_id = "nielsr/image-segmentation-toy-data"
        image1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_image_1.png", repo_type="dataset")
        )
        image2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_image_2.png", repo_type="dataset")
        )
        annotation1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_annotation_1.png", repo_type="dataset")
        )
        annotation2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_annotation_2.png", repo_type="dataset")
        )

        # get instance segmentations and instance-to-segmentation mappings
        def get_instance_segmentation_and_mapping(annotation):
            instance_seg = np.array(annotation)[:, :, 1]
            class_id_map = np.array(annotation)[:, :, 0]
            class_labels = np.unique(class_id_map)

            # create mapping between instance IDs and semantic category IDs
            inst2class = {}
            for label in class_labels:
                instance_ids = np.unique(instance_seg[class_id_map == label])
                inst2class.update({i: label for i in instance_ids})

            return instance_seg, inst2class

        instance_seg1, inst2class1 = get_instance_segmentation_and_mapping(annotation1)
        instance_seg2, inst2class2 = get_instance_segmentation_and_mapping(annotation2)

        # create a feature extractor
        feature_extractor = MaskFormerImageProcessor(reduce_labels=True, ignore_index=255, size=(512, 512))

        # prepare the images and annotations
        inputs = feature_extractor(
            [image1, image2],
            [instance_seg1, instance_seg2],
            instance_id_to_semantic_id=[inst2class1, inst2class2],
            return_tensors="pt",
        )

        # verify the pixel values and pixel mask
        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))

        # verify the class labels
        self.assertEqual(len(inputs["class_labels"]), 2)
        self.assertTrue(torch.allclose(inputs["class_labels"][0], torch.tensor([30, 55])))
        self.assertTrue(torch.allclose(inputs["class_labels"][1], torch.tensor([4, 4, 23, 55])))

        # verify the mask labels
        self.assertEqual(len(inputs["mask_labels"]), 2)
        self.assertEqual(inputs["mask_labels"][0].shape, (2, 512, 512))
        self.assertEqual(inputs["mask_labels"][1].shape, (4, 512, 512))
        self.assertEquals(inputs["mask_labels"][0].sum().item(), 41527.0)
        self.assertEquals(inputs["mask_labels"][1].sum().item(), 26259.0)
//  [python]
def test_integration_semantic_segmentation(self):
        # load 2 images and corresponding semantic annotations from the hub
        repo_id = "nielsr/image-segmentation-toy-data"
        image1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="semantic_segmentation_image_1.png", repo_type="dataset")
        )
        image2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="semantic_segmentation_image_2.png", repo_type="dataset")
        )
        annotation1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="semantic_segmentation_annotation_1.png", repo_type="dataset")
        )
        annotation2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="semantic_segmentation_annotation_2.png", repo_type="dataset")
        )

        # create a feature extractor
        feature_extractor = MaskFormerFeatureExtractor(reduce_labels=True, ignore_index=255, size=(512, 512))

        # prepare the images and annotations
        inputs = feature_extractor(
            [image1, image2],
            [annotation1, annotation2],
            return_tensors="pt",
        )

        # verify the pixel values and pixel mask
        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))

        # verify the class labels
        self.assertEqual(len(inputs["class_labels"]), 2)
        self.assertTrue(torch.allclose(inputs["class_labels"][0], torch.tensor([2, 4, 60])))
        self.assertTrue(torch.allclose(inputs["class_labels"][1], torch.tensor([0, 3, 7, 8, 15, 28, 30, 143])))

        # verify the mask labels
        self.assertEqual(len(inputs["mask_labels"]), 2)
        self.assertEqual(inputs["mask_labels"][0].shape, (3, 512, 512))
        self.assertEqual(inputs["mask_labels"][1].shape, (8, 512, 512))
        self.assertEquals(inputs["mask_labels"][0].sum().item(), 170200.0)
        self.assertEquals(inputs["mask_labels"][1].sum().item(), 257036.0)
//  [python]
def test_integration_semantic_segmentation(self):
        # load 2 images and corresponding semantic annotations from the hub
        repo_id = "nielsr/image-segmentation-toy-data"
        image1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="semantic_segmentation_image_1.png", repo_type="dataset")
        )
        image2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="semantic_segmentation_image_2.png", repo_type="dataset")
        )
        annotation1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="semantic_segmentation_annotation_1.png", repo_type="dataset")
        )
        annotation2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="semantic_segmentation_annotation_2.png", repo_type="dataset")
        )

        # create a feature extractor
        feature_extractor = MaskFormerImageProcessor(reduce_labels=True, ignore_index=255, size=(512, 512))

        # prepare the images and annotations
        inputs = feature_extractor(
            [image1, image2],
            [annotation1, annotation2],
            return_tensors="pt",
        )

        # verify the pixel values and pixel mask
        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))

        # verify the class labels
        self.assertEqual(len(inputs["class_labels"]), 2)
        self.assertTrue(torch.allclose(inputs["class_labels"][0], torch.tensor([2, 4, 60])))
        self.assertTrue(torch.allclose(inputs["class_labels"][1], torch.tensor([0, 3, 7, 8, 15, 28, 30, 143])))

        # verify the mask labels
        self.assertEqual(len(inputs["mask_labels"]), 2)
        self.assertEqual(inputs["mask_labels"][0].shape, (3, 512, 512))
        self.assertEqual(inputs["mask_labels"][1].shape, (8, 512, 512))
        self.assertEquals(inputs["mask_labels"][0].sum().item(), 170200.0)
        self.assertEquals(inputs["mask_labels"][1].sum().item(), 257036.0)
//  [python]
def test_integration_instance_segmentation(self):
        # load 2 images and corresponding annotations from the hub
        repo_id = "nielsr/image-segmentation-toy-data"
        image1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_image_1.png", repo_type="dataset")
        )
        image2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_image_2.png", repo_type="dataset")
        )
        annotation1 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_annotation_1.png", repo_type="dataset")
        )
        annotation2 = Image.open(
            hf_hub_download(repo_id=repo_id, filename="instance_segmentation_annotation_2.png", repo_type="dataset")
        )

        # get instance segmentations and instance-to-segmentation mappings
        def get_instance_segmentation_and_mapping(annotation):
            instance_seg = np.array(annotation)[:, :, 1]
            class_id_map = np.array(annotation)[:, :, 0]
            class_labels = np.unique(class_id_map)

            # create mapping between instance IDs and semantic category IDs
            inst2class = {}
            for label in class_labels:
                instance_ids = np.unique(instance_seg[class_id_map == label])
                inst2class.update({i: label for i in instance_ids})

            return instance_seg, inst2class

        instance_seg1, inst2class1 = get_instance_segmentation_and_mapping(annotation1)
        instance_seg2, inst2class2 = get_instance_segmentation_and_mapping(annotation2)

        # create a image processor
        image_processing = MaskFormerFeatureExtractor(reduce_labels=True, ignore_index=255, size=(512, 512))

        # prepare the images and annotations
        inputs = image_processing(
            [image1, image2],
            [instance_seg1, instance_seg2],
            instance_id_to_semantic_id=[inst2class1, inst2class2],
            return_tensors="pt",
        )

        # verify the pixel values and pixel mask
        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))

        # verify the class labels
        self.assertEqual(len(inputs["class_labels"]), 2)
        self.assertTrue(torch.allclose(inputs["class_labels"][0], torch.tensor([30, 55])))
        self.assertTrue(torch.allclose(inputs["class_labels"][1], torch.tensor([4, 4, 23, 55])))

        # verify the mask labels
        self.assertEqual(len(inputs["mask_labels"]), 2)
        self.assertEqual(inputs["mask_labels"][0].shape, (2, 512, 512))
        self.assertEqual(inputs["mask_labels"][1].shape, (4, 512, 512))
        self.assertEquals(inputs["mask_labels"][0].sum().item(), 41527.0)
        self.assertEquals(inputs["mask_labels"][1].sum().item(), 26259.0)