diff --git a/google/generativeai/__init__.py b/google/generativeai/__init__.py index 73025a1b4..66af17641 100644 --- a/google/generativeai/__init__.py +++ b/google/generativeai/__init__.py @@ -79,7 +79,6 @@ __version__ = version.__version__ -del embedding del files del generative_models del models diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py index f3db610e1..23241a536 100644 --- a/google/generativeai/types/content_types.py +++ b/google/generativeai/types/content_types.py @@ -35,6 +35,7 @@ import IPython.display IMAGE_TYPES = (PIL.Image.Image, IPython.display.Image) + ImageType = PIL.Image.Image | IPython.display.Image else: IMAGE_TYPES = () try: @@ -52,6 +53,8 @@ except ImportError: IPython = None + ImageType = Union["PIL.Image.Image", "IPython.display.Image"] + __all__ = [ "BlobDict", @@ -123,7 +126,7 @@ def webp_blob(image: PIL.Image.Image) -> protos.Blob: return file_blob(image) or webp_blob(image) -def image_to_blob(image) -> protos.Blob: +def image_to_blob(image: ImageType) -> protos.Blob: if PIL is not None: if isinstance(image, PIL.Image.Image): return _pil_to_blob(image) diff --git a/google/generativeai/vision_models/__init__.py b/google/generativeai/vision_models/__init__.py index 65a545831..2a4a27e32 100644 --- a/google/generativeai/vision_models/__init__.py +++ b/google/generativeai/vision_models/__init__.py @@ -15,6 +15,7 @@ """Classes for working with vision models.""" from google.generativeai.vision_models._vision_models import ( + check_watermark, Image, GeneratedImage, ImageGenerationModel, diff --git a/google/generativeai/vision_models/_vision_models.py b/google/generativeai/vision_models/_vision_models.py index 3e7dbfa22..52ec689a9 100644 --- a/google/generativeai/vision_models/_vision_models.py +++ b/google/generativeai/vision_models/_vision_models.py @@ -18,15 +18,16 @@ import base64 import collections import dataclasses -import hashlib import io import json +import os import pathlib import typing from typing import Any, Dict, List, Literal, Optional, Union from google.generativeai import client from google.generativeai import protos +from google.generativeai.types import content_types from google.protobuf import struct_pb2 @@ -110,6 +111,52 @@ def to_mapping_value(value) -> struct_pb2.Struct: PersonGeneration = Literal["dont_allow", "allow_adult"] PERSON_GENERATIONS = PersonGeneration.__args__ # type: ignore +ImageLikeType = Union["Image", pathlib.Path, content_types.ImageType] + + +def check_watermark( + img: ImageLikeType, model_id: str = "models/image-verification-001" +) -> "CheckWatermarkResult": + """Checks if an image has a Google-AI watermark. + + Args: + img: can be a `pathlib.Path` or a `PIL.Image.Image`, `IPython.display.Image`, or `google.generativeai.Image`. + model_id: Which version of the image-verification model to send the image to. + + Returns: + + """ + if isinstance(img, Image): + pass + elif isinstance(img, pathlib.Path): + img = Image.load_from_file(img) + elif IPython_display is not None and isinstance(img, IPython_display.Image): + img = Image(image_bytes=img.data) + elif PIL_Image is not None and isinstance(img, PIL_Image.Image): + blob = content_types._pil_to_blob(img) + img = Image(image_bytes=blob.data) + elif isinstance(img, protos.Blob): + img = Image(image_bytes=img.data) + else: + raise TypeError( + f"Not implemented: Could not convert a {type(img)} into `Image`\n {img=}" + ) + + prediction_client = client.get_default_prediction_client() + if not model_id.startswith("models/"): + model_id = f"models/{model_id}" + + instance = {"image": {"bytesBase64Encoded": base64.b64encode(img._loaded_bytes).decode()}} + parameters = {"watermarkVerification": True} + + # This is to get around https://github.com/googleapis/proto-plus-python/issues/488 + pr = protos.PredictRequest.pb() + request = pr(model=model_id, instances=[to_value(instance)], parameters=to_value(parameters)) + + response = prediction_client.predict(request) + + return CheckWatermarkResult(response.predictions) + class Image: """Image.""" @@ -131,7 +178,7 @@ def __init__( self._image_bytes = image_bytes @staticmethod - def load_from_file(location: str) -> "Image": + def load_from_file(location: os.PathLike) -> "Image": """Loads image from local file or Google Cloud Storage. Args: @@ -206,6 +253,29 @@ def _as_base64_string(self) -> str: def _repr_png_(self): return self._pil_image._repr_png_() # type:ignore + check_watermark = check_watermark + + +class CheckWatermarkResult: + def __init__(self, predictions): + self._predictions = predictions + + @property + def decision(self): + return self._predictions[0]["decision"] + + def __str__(self): + return f"CheckWatermarkResult([{{'decision': {self.decision!r}}}])" + + def __bool__(self): + decision = self.decision + if decision == "ACCEPT": + return True + elif decision == "REJECT": + return False + else: + raise ValueError(f"Unrecognized result: {decision}") + class ImageGenerationModel: """Generates images from text prompt. @@ -479,7 +549,7 @@ def generation_parameters(self): return self._generation_parameters @staticmethod - def load_from_file(location: str) -> "GeneratedImage": + def load_from_file(location: os.PathLike) -> "GeneratedImage": """Loads image from file. Args: