feat(Hugging face): Vision methods - image classification / image segmentation / object detection

examples/hugging-face/image-to-image-async.py

@@ -1,18 +1,19 @@
 import asyncio
+import argparse
 
 from huggingface_hub import AsyncInferenceClient
 from scope3ai import Scope3AI
 from scope3ai.tracers.huggingface.image_to_image import HUGGING_FACE_IMAGE_TO_IMAGE_TASK
 
 
-async def main():
+async def main(image_path, model):
     client = AsyncInferenceClient()
     scope3 = Scope3AI.init()
-    model = await client.get_recommended_model(HUGGING_FACE_IMAGE_TO_IMAGE_TASK)
+    if not model:
+        model = await client.get_recommended_model(HUGGING_FACE_IMAGE_TO_IMAGE_TASK)
 
     with scope3.trace() as tracer:
         # Replace `image_path` with the path to your image file
-        image_path = "path/to/your/image.jpg"
         with open(image_path, "rb") as f:
             response = await client.image_to_image(
                 model=model,
@@ -28,4 +29,16 @@ async def main():
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    parser = argparse.ArgumentParser(
+        description="Run image-to-image transformation with environmental impact tracing."
+    )
+    parser.add_argument("--image-path", type=str, help="Path to the input image file.")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Hugging Face model to use. Defaults to the recommended model.",
+    )
+    args = parser.parse_args()
+
+    asyncio.run(main(args.image_path, args.model))

examples/hugging-face/image-to-image.py

@@ -1,18 +1,20 @@
 import asyncio
+import argparse
 
 from huggingface_hub import InferenceClient
 from scope3ai import Scope3AI
 from scope3ai.tracers.huggingface.image_to_image import HUGGING_FACE_IMAGE_TO_IMAGE_TASK
 
 
-async def main():
+async def main(image_path, model):
     client = InferenceClient()
     scope3 = Scope3AI.init()
-    model = client.get_recommended_model(HUGGING_FACE_IMAGE_TO_IMAGE_TASK)
+
+    # Use recommended model if none is provided
+    if not model:
+        model = client.get_recommended_model(HUGGING_FACE_IMAGE_TO_IMAGE_TASK)
 
     with scope3.trace() as tracer:
-        # Replace `image_path` with the path to your image file
-        image_path = "path/to/your/image.jpg"
         with open(image_path, "rb") as f:
             response = client.image_to_image(
                 model=model,
@@ -28,4 +30,16 @@ async def main():
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    parser = argparse.ArgumentParser(
+        description="Run image-to-image transformation with environmental impact tracing."
+    )
+    parser.add_argument("--image-path", type=str, help="Path to the input image file.")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Hugging Face model to use. Defaults to the recommended model.",
+    )
+    args = parser.parse_args()
+
+    asyncio.run(main(args.image_path, args.model))

examples/hugging-face/speech-to-text-async.py

@@ -1,18 +1,19 @@
+import argparse
 import asyncio
 
 from huggingface_hub import AsyncInferenceClient
 from scope3ai import Scope3AI
 from scope3ai.tracers.huggingface.speech_to_text import HUGGING_FACE_SPEECH_TO_TEXT_TASK
 
 
-async def main():
+async def main(audio_path, model):
     client = AsyncInferenceClient()
     scope3 = Scope3AI.init()
-    model = await client.get_recommended_model(HUGGING_FACE_SPEECH_TO_TEXT_TASK)
+    if not model:
+        model = await client.get_recommended_model(HUGGING_FACE_SPEECH_TO_TEXT_TASK)
 
     with scope3.trace() as tracer:
         # Replace `audio_path` with the path to your audio file
-        audio_path = "path/to/your/audio.wav"
         with open(audio_path, "rb") as f:
             response = await client.automatic_speech_recognition(
                 model=model,
@@ -27,4 +28,16 @@ async def main():
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    parser = argparse.ArgumentParser(
+        description="Run speech-to-text transformation with environmental impact tracing."
+    )
+    parser.add_argument("--audio_path", type=str, help="Path to the input audio file.")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Hugging Face model to use. Defaults to the recommended model.",
+    )
+    args = parser.parse_args()
+
+    asyncio.run(main(args.audio_path, args.model))

examples/hugging-face/speech-to-text.py

@@ -1,18 +1,19 @@
+import argparse
 import asyncio
 
 from huggingface_hub import InferenceClient
 from scope3ai import Scope3AI
 from scope3ai.tracers.huggingface.speech_to_text import HUGGING_FACE_SPEECH_TO_TEXT_TASK
 
 
-async def main():
+async def main(audio_path, model):
     client = InferenceClient()
     scope3 = Scope3AI.init()
-    model = client.get_recommended_model(HUGGING_FACE_SPEECH_TO_TEXT_TASK)
+    if not model:
+        model = client.get_recommended_model(HUGGING_FACE_SPEECH_TO_TEXT_TASK)
 
     with scope3.trace() as tracer:
         # Replace `audio_path` with the path to your audio file
-        audio_path = "path/to/your/audio.wav"
         with open(audio_path, "rb") as f:
             response = client.automatic_speech_recognition(
                 model=model,
@@ -27,4 +28,16 @@ async def main():
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    parser = argparse.ArgumentParser(
+        description="Run speech-to-text transformation with environmental impact tracing."
+    )
+    parser.add_argument("--audio_path", type=str, help="Path to the input audio file.")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Hugging Face model to use. Defaults to the recommended model.",
+    )
+    args = parser.parse_args()
+
+    asyncio.run(main(args.audio_path, args.model))

scope3ai/tracers/huggingface/image_to_image.py

@@ -1,3 +1,4 @@
+import io
 import time
 from dataclasses import dataclass
 from typing import Any, Callable, Optional, Union
@@ -41,9 +42,12 @@ def _hugging_face_image_to_image_wrapper(
         encoder = tiktoken.get_encoding("cl100k_base")
         prompt = args[1] if len(args) > 1 else kwargs.get("prompt", "")
         input_tokens = len(encoder.encode(prompt)) if prompt != "" else 0
-    input_images = None
     try:
-        input_image = Image.open(args[0] if len(args) > 0 else kwargs["image"])
+        image_param = args[0] if len(args) > 0 else kwargs["image"]
+        if type(image_param) is str:
+            input_image = Image.open(args[0] if len(args) > 0 else kwargs["image"])
+        else:
+            input_image = Image.open(io.BytesIO(image_param))
         input_width, input_height = input_image.size
         input_images = [
             ("{width}x{height}".format(width=input_width, height=input_height))

scope3ai/tracers/huggingface/instrument.py

@@ -23,6 +23,18 @@
     huggingface_translation_wrapper_non_stream,
     huggingface_translation_wrapper_async_non_stream,
 )
+from scope3ai.tracers.huggingface.vision.image_classification import (
+    huggingface_image_classification_wrapper,
+    huggingface_image_classification_wrapper_async,
+)
+from scope3ai.tracers.huggingface.vision.image_segmentation import (
+    huggingface_image_segmentation_wrapper,
+    huggingface_image_segmentation_wrapper_async,
+)
+from scope3ai.tracers.huggingface.vision.object_detection import (
+    huggingface_object_detection_wrapper,
+    huggingface_object_detection_wrapper_async,
+)
 
 
 class HuggingfaceInstrumentor:
@@ -83,6 +95,36 @@ def __init__(self) -> None:
                 "name": "AsyncInferenceClient.image_to_image",
                 "wrapper": huggingface_image_to_image_wrapper_async,
             },
+            {
+                "module": "huggingface_hub.inference._client",
+                "name": "InferenceClient.image_classification",
+                "wrapper": huggingface_image_classification_wrapper,
+            },
+            {
+                "module": "huggingface_hub.inference._generated._async_client",
+                "name": "AsyncInferenceClient.image_classification",
+                "wrapper": huggingface_image_classification_wrapper_async,
+            },
+            {
+                "module": "huggingface_hub.inference._client",
+                "name": "InferenceClient.image_segmentation",
+                "wrapper": huggingface_image_segmentation_wrapper,
+            },
+            {
+                "module": "huggingface_hub.inference._generated._async_client",
+                "name": "AsyncInferenceClient.image_segmentation",
+                "wrapper": huggingface_image_segmentation_wrapper_async,
+            },
+            {
+                "module": "huggingface_hub.inference._client",
+                "name": "InferenceClient.object_detection",
+                "wrapper": huggingface_object_detection_wrapper,
+            },
+            {
+                "module": "huggingface_hub.inference._generated._async_client",
+                "name": "AsyncInferenceClient.object_detection",
+                "wrapper": huggingface_object_detection_wrapper_async,
+            },
         ]
 
     def instrument(self) -> None:

scope3ai/tracers/huggingface/vision/image_classification.py

@@ -0,0 +1,106 @@
+import io
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Union, List
+from PIL import Image
+from aiohttp import ClientResponse
+from huggingface_hub import (
+    InferenceClient,
+    AsyncInferenceClient,
+    ImageClassificationOutputElement,
+)  # type: ignore[import-untyped]
+from requests import Response
+
+from scope3ai.api.types import Scope3AIContext, Model, ImpactRow
+from scope3ai.api.typesgen import Task
+from scope3ai.constants import PROVIDERS
+from scope3ai.lib import Scope3AI
+from scope3ai.response_interceptor.aiohttp_interceptor import aiohttp_response_capture
+from scope3ai.response_interceptor.requests_interceptor import requests_response_capture
+
+PROVIDER = PROVIDERS.HUGGINGFACE_HUB.value
+HUGGING_FACE_IMAGE_CLASSIFICATION_TASK = "image-classification"
+
+
+@dataclass
+class ImageClassificationOutput:
+    elements: List[ImageClassificationOutputElement] = None
+    scope3ai: Optional[Scope3AIContext] = None
+
+
+def _hugging_face_image_classification_wrapper(
+    timer_start: Any,
+    model: Any,
+    response: Any,
+    http_response: Union[ClientResponse, Response],
+    args: Any,
+    kwargs: Any,
+) -> ImageClassificationOutput:
+    input_tokens = 0
+    if http_response:
+        compute_time = http_response.headers.get("x-compute-time")
+    else:
+        compute_time = time.perf_counter() - timer_start
+    try:
+        image_param = args[0] if len(args) > 0 else kwargs["image"]
+        if type(image_param) is str:
+            input_image = Image.open(args[0] if len(args) > 0 else kwargs["image"])
+        else:
+            input_image = Image.open(io.BytesIO(image_param))
+        input_width, input_height = input_image.size
+        input_images = [
+            ("{width}x{height}".format(width=input_width, height=input_height))
+        ]
+    except Exception:
+        pass
+    scope3_row = ImpactRow(
+        model=Model(id=model),
+        input_tokens=input_tokens,
+        task=Task.image_classification,
+        output_images=[],  # No images to output in classification
+        request_duration_ms=float(compute_time) * 1000,
+        managed_service_id=PROVIDER,
+        input_images=input_images,
+    )
+
+    scope3_ctx = Scope3AI.get_instance().submit_impact(scope3_row)
+    result = ImageClassificationOutput()
+    result.elements = response
+    result.scope3ai = scope3_ctx
+    return result
+
+
+def huggingface_image_classification_wrapper(
+    wrapped: Callable, instance: InferenceClient, args: Any, kwargs: Any
+) -> ImageClassificationOutput:
+    timer_start = time.perf_counter()
+    http_response: Response | None = None
+    with requests_response_capture() as responses:
+        response = wrapped(*args, **kwargs)
+        http_responses = responses.get()
+        if len(http_responses) > 0:
+            http_response = http_responses[-1]
+    model = kwargs.get("model") or instance.get_recommended_model(
+        HUGGING_FACE_IMAGE_CLASSIFICATION_TASK
+    )
+    return _hugging_face_image_classification_wrapper(
+        timer_start, model, response, http_response, args, kwargs
+    )
+
+
+async def huggingface_image_classification_wrapper_async(
+    wrapped: Callable, instance: AsyncInferenceClient, args: Any, kwargs: Any
+) -> ImageClassificationOutput:
+    timer_start = time.perf_counter()
+    http_response: ClientResponse | None = None
+    with aiohttp_response_capture() as responses:
+        response = await wrapped(*args, **kwargs)
+        http_responses = responses.get()
+        if len(http_responses) > 0:
+            http_response = http_responses[-1]
+    model = kwargs.get("model") or instance.get_recommended_model(
+        HUGGING_FACE_IMAGE_CLASSIFICATION_TASK
+    )
+    return _hugging_face_image_classification_wrapper(
+        timer_start, model, response, http_response, args, kwargs
+    )