Skip to content

Commit

Permalink
Remove watermark
Browse files Browse the repository at this point in the history
Change-Id: Ie8f26fb7c4eb26c4cf7bdd746b608a652dfbd5a6
  • Loading branch information
MarkDaoust committed Nov 12, 2024
1 parent c72aa6b commit 302ccc7
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 44 deletions.
43 changes: 1 addition & 42 deletions google/generativeai/types/image_types/_image_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
from typing import Any, Dict, Optional, Union

from google.generativeai import protos
from google.generativeai import client

# pylint: disable=g-import-not-at-top
if typing.TYPE_CHECKING:
Expand Down Expand Up @@ -38,7 +37,7 @@
ImageType = Union["Image", "PIL.Image.Image", "IPython.display.Image"]
# pylint: enable=g-import-not-at-top

__all__ = ["Image", "GeneratedImage", "check_watermark", "CheckWatermarkResult", "ImageType"]
__all__ = ["Image", "GeneratedImage", "ImageType"]


def _pil_to_blob(image: PIL.Image.Image) -> protos.Blob:
Expand Down Expand Up @@ -121,46 +120,6 @@ def __bool__(self):
raise ValueError(f"Unrecognized result: {decision}")


def check_watermark(
img: Union[pathlib.Path, ImageType], 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 = _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}

response = prediction_client.predict(
model=model_id, instances=[instance], parameters=parameters
)

return CheckWatermarkResult(response.predictions)


class Image:
Expand Down
3 changes: 1 addition & 2 deletions google/generativeai/vision_models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,14 @@
#
"""Classes for working with vision models."""

from google.generativeai.types.image_types import check_watermark, Image, GeneratedImage
from google.generativeai.types.image_types import Image, GeneratedImage

from google.generativeai.vision_models._vision_models import (
ImageGenerationModel,
ImageGenerationResponse,
)

__all__ = [
"check_watermark",
"Image",
"GeneratedImage",
"ImageGenerationModel",
Expand Down

0 comments on commit 302ccc7

Please sign in to comment.