|
from PIL import Image, ImageDraw, ImageFilter, ImageFont |
|
import numpy as np |
|
from transformers.tools import Tool |
|
|
|
class PlaceholderImageTool(Tool): |
|
"""Replacement 'image_generator' |
|
- transformers HfAgent likes to use the image_generator tool. |
|
- I don't have the disk space for that. |
|
- This replacement provides a random placeholder image instead. |
|
Install it thusly: `agent.toolbox['image_generator']=PlaceholderImageTool()` |
|
""" |
|
name = "image_generator" |
|
description = "I will generate an image based on a prompt." |
|
inputs = ["text"] |
|
outputs = ["image"] |
|
def sepia(self, dim:int=64) -> Image: |
|
"produce a sepia image, size=(dim,dim,3)" |
|
data = np.random.randint(0,255,(dim,dim,3), dtype=np.uint8) |
|
image = Image.fromarray(data) |
|
sepia_matrix = [ |
|
0.393, 0.769, 0.189, 0, |
|
0.349, 0.686, 0.168, 0, |
|
0.272, 0.534, 0.131, 0 |
|
] |
|
applied = image.convert('RGB', sepia_matrix) |
|
return applied |
|
def __call__(self, prompt): |
|
"returns a sepia image overlaid with a random digit (1,6 inclusive)" |
|
DIM=64 |
|
image = self.sepia(DIM) |
|
|
|
draw = ImageDraw.Draw(image) |
|
digit = str(np.random.randint(1,7)) |
|
font = ImageFont.truetype("./Lora-Bold.ttf", size=48) |
|
draw.text((DIM//2,DIM//2), digit, fill=(0,0,0), font=font, anchor="mm") |
|
return image |
