fix(pipelines): respect problem_type in image and video classification pipelines (#50582)

ImageClassificationPipeline checked problem_type == multi_label_classification
or num_labels == 1 before checking problem_type == regression, so regression
models had their outputs squashed through sigmoid (for num_labels == 1) or
softmax (for num_labels > 1). Check problem_type == regression first, matching
TextClassificationPipeline (#34066).

VideoClassificationPipeline unconditionally defaulted function_to_apply to
softmax in _sanitize_parameters(), ignoring model.config.problem_type and
overwriting initialization-time function_to_apply overrides on subsequent calls.
Defer the default to postprocess() based on problem_type, matching
AudioClassificationPipeline (#50329).

Co-authored-by: surendra902 <surendra902@users.noreply.github.com>

Author: surendra902

src/transformers/pipelines/image_classification.py

@@ -65,7 +65,8 @@ class ClassificationFunction(ExplicitEnum):
             The function to apply to the model outputs in order to retrieve the scores. Accepts four different values:
 
             - `"default"`: if the model has a single label, will apply the sigmoid function on the output. If the model
-              has several labels, will apply the softmax function on the output.
+              has several labels, will apply the softmax function on the output. In case of regression tasks, will not
+              apply any function on the output.
             - `"sigmoid"`: Applies the sigmoid function on the output.
             - `"softmax"`: Applies the softmax function on the output.
             - `"none"`: Does not apply any function on the output.""",
@@ -148,6 +149,7 @@ def __call__(
                 If this argument is not specified, then it will apply the following functions according to the number
                 of labels:
 
+                - If problem type is regression, will not apply any function on the output.
                 - If the model has a single label, will apply the sigmoid function on the output.
                 - If the model has several labels, will apply the softmax function on the output.
 
@@ -193,7 +195,9 @@ def _forward(self, model_inputs):
 
     def postprocess(self, model_outputs, function_to_apply=None, top_k=5):
         if function_to_apply is None:
-            if self.model.config.problem_type == "multi_label_classification" or self.model.config.num_labels == 1:
+            if self.model.config.problem_type == "regression":
+                function_to_apply = ClassificationFunction.NONE
+            elif self.model.config.problem_type == "multi_label_classification" or self.model.config.num_labels == 1:
                 function_to_apply = ClassificationFunction.SIGMOID
             elif self.model.config.problem_type == "single_label_classification" or self.model.config.num_labels > 1:
                 function_to_apply = ClassificationFunction.SOFTMAX

src/transformers/pipelines/video_classification.py

@@ -92,8 +92,6 @@ def _sanitize_parameters(self, top_k=None, num_frames=None, frame_sampling_rate=
                     "Valid options are ['softmax', 'sigmoid', 'none']"
                 )
             postprocess_params["function_to_apply"] = function_to_apply
-        else:
-            postprocess_params["function_to_apply"] = "softmax"
         return preprocess_params, {}, postprocess_params
 
     @overload
@@ -124,11 +122,11 @@ def __call__(self, inputs: str | list[str] | None, **kwargs):
             frame_sampling_rate (`int`, *optional*, defaults to 1):
                 The sampling rate used to select frames from the video. If not provided, will default to 1, i.e. every
                 frame will be used.
-            function_to_apply(`str`, *optional*, defaults to "softmax"):
-                The function to apply to the model output. By default, the pipeline will apply the softmax function to
-                the output of the model. Valid options: ["softmax", "sigmoid", "none"]. Note that passing Python's
-                built-in `None` will default to "softmax", so you need to pass the string "none" to disable any
-                post-processing.
+            function_to_apply (`str`, *optional*):
+                The function to apply to the model output. Defaults to "none" for regression, "sigmoid" for
+                multi-label classification, and "softmax" otherwise. Valid options: ["softmax", "sigmoid", "none"].
+                Explicit values supplied at initialization or call time override these defaults.
+                Pass the string "none" to disable post-processing.
 
         Return:
             A list of dictionaries or a list of list of dictionaries containing result. If the input is a single video,
@@ -170,7 +168,15 @@ def _forward(self, model_inputs):
         model_outputs = self.model(**model_inputs)
         return model_outputs
 
-    def postprocess(self, model_outputs, top_k=5, function_to_apply="softmax"):
+    def postprocess(self, model_outputs, top_k=5, function_to_apply=None):
+        if function_to_apply is None:
+            if self.model.config.problem_type == "regression":
+                function_to_apply = "none"
+            elif self.model.config.problem_type == "multi_label_classification":
+                function_to_apply = "sigmoid"
+            else:
+                function_to_apply = "softmax"
+
         if top_k > self.model.config.num_labels:
             top_k = self.model.config.num_labels
 

tests/pipelines/test_pipelines_image_classification.py

@@ -173,6 +173,13 @@ def test_small_model_pt(self):
             ],
         )
 
+        image_classifier.model.config.problem_type = "regression"
+        outputs = image_classifier("./tests/fixtures/tests_samples/COCO/000000039769.png")
+        self.assertEqual(
+            nested_simplify(outputs, decimals=4),
+            [{"label": "LABEL_1", "score": 0.1426}, {"label": "LABEL_0", "score": -0.1557}],
+        )
+
     def test_custom_tokenizer(self):
         tokenizer = PreTrainedTokenizerBase()
 

tests/pipelines/test_pipelines_video_classification.py

@@ -16,7 +16,12 @@
 
 from huggingface_hub import VideoClassificationOutputElement, hf_hub_download
 
-from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VideoMAEImageProcessor
+from transformers import (
+    MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
+    AutoModelForVideoClassification,
+    VideoMAEImageProcessor,
+    is_torch_available,
+)
 from transformers.pipelines import VideoClassificationPipeline, pipeline
 from transformers.testing_utils import (
     compare_pipeline_output_to_hub_spec,
@@ -30,6 +35,10 @@
 from .test_pipelines_common import ANY
 
 
+if is_torch_available():
+    import torch
+
+
 @is_pipeline_test
 @require_torch
 @require_vision
@@ -122,3 +131,64 @@ def test_small_model_pt(self):
         for output in outputs:
             for element in output:
                 compare_pipeline_output_to_hub_spec(element, VideoClassificationOutputElement)
+
+    @require_torch
+    def test_problem_type_defaults_and_overrides(self):
+        small_model = AutoModelForVideoClassification.from_pretrained(
+            "hf-internal-testing/tiny-random-VideoMAEForVideoClassification"
+        )
+        small_image_processor = VideoMAEImageProcessor(
+            size={"shortest_edge": 10}, crop_size={"height": 10, "width": 10}
+        )
+        classifier = pipeline(
+            "video-classification",
+            model=small_model,
+            image_processor=small_image_processor,
+            frame_sampling_rate=4,
+        )
+        video_file_path = hf_hub_download(repo_id="nateraw/video-demo", filename="archery.mp4", repo_type="dataset")
+        num_labels = classifier.model.config.num_labels
+        labels = [classifier.model.config.id2label[i] for i in range(num_labels)]
+
+        raw = classifier(video_file_path, top_k=num_labels, function_to_apply="none")
+        raw_scores = {item["label"]: item["score"] for item in raw}
+        logits = torch.tensor([raw_scores[label] for label in labels])
+
+        for problem_type, activation in (
+            ("regression", "none"),
+            ("multi_label_classification", "sigmoid"),
+            ("single_label_classification", "softmax"),
+            (None, "softmax"),
+        ):
+            with self.subTest(problem_type=problem_type):
+                classifier.model.config.problem_type = problem_type
+                result = classifier(video_file_path, top_k=num_labels)
+                scores = {item["label"]: item["score"] for item in result}
+                expected = {
+                    "none": logits,
+                    "sigmoid": logits.sigmoid(),
+                    "softmax": logits.softmax(-1),
+                }[activation]
+                self.assertEqual(len(scores), num_labels)
+                for label, value in zip(labels, expected.tolist()):
+                    self.assertAlmostEqual(scores[label], value, places=4)
+
+        classifier.model.config.problem_type = "regression"
+        initialized = pipeline(
+            "video-classification",
+            model=classifier.model,
+            image_processor=small_image_processor,
+            frame_sampling_rate=4,
+            function_to_apply="softmax",
+        )
+
+        for override, expected in (
+            (None, logits.softmax(-1)),
+            ("none", logits),
+            ("sigmoid", logits.sigmoid()),
+        ):
+            with self.subTest(call_override=override):
+                result = initialized(video_file_path, top_k=num_labels, function_to_apply=override)
+                scores = {item["label"]: item["score"] for item in result}
+                for label, value in zip(labels, expected.tolist()):
+                    self.assertAlmostEqual(scores[label], value, places=4)