File size: 6,923 Bytes
023c371 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
from PIL import Image, ImageDraw
import io
import random
import math
class EndpointHandler:
def __init__(self):
"""Initialize the DiffSketcher handler with fallback PIL drawing"""
pass
def __call__(self, data):
"""
Generate a sketch-style image using PIL drawing (fallback method)
Args:
data (dict): Input data containing:
- inputs (str): Text prompt
- parameters (dict): Generation parameters
Returns:
PIL.Image.Image: Generated sketch image
"""
try:
# Extract inputs
prompt = data.get("inputs", "")
parameters = data.get("parameters", {})
# Extract parameters
width = parameters.get("width", 224)
height = parameters.get("height", 224)
guidance_scale = parameters.get("guidance_scale", 7.5)
seed = parameters.get("seed", 42)
# Set random seed for reproducibility
random.seed(seed)
# Create white background
image = Image.new('RGB', (width, height), 'white')
draw = ImageDraw.Draw(image)
# Generate sketch based on prompt keywords
self._draw_sketch_from_prompt(draw, prompt, width, height)
return image
except Exception as e:
# Return error image
error_image = Image.new('RGB', (224, 224), 'white')
error_draw = ImageDraw.Draw(error_image)
error_draw.text((10, 100), f"Error: {str(e)[:30]}", fill='red')
return error_image
def _draw_sketch_from_prompt(self, draw, prompt, width, height):
"""Draw a simple sketch based on prompt keywords"""
prompt_lower = prompt.lower()
# Define colors for sketching
colors = ['black', 'gray', 'darkgray']
if any(word in prompt_lower for word in ['car', 'vehicle', 'automobile']):
self._draw_car(draw, width, height, colors)
elif any(word in prompt_lower for word in ['house', 'building', 'home']):
self._draw_house(draw, width, height, colors)
elif any(word in prompt_lower for word in ['flower', 'plant', 'bloom']):
self._draw_flower(draw, width, height, colors)
elif any(word in prompt_lower for word in ['tree', 'forest']):
self._draw_tree(draw, width, height, colors)
elif any(word in prompt_lower for word in ['mountain', 'landscape']):
self._draw_mountain(draw, width, height, colors)
else:
self._draw_abstract(draw, width, height, colors)
def _draw_car(self, draw, width, height, colors):
"""Draw a simple car sketch"""
cx, cy = width // 2, height // 2
# Car body
draw.rectangle([cx-60, cy-20, cx+60, cy+20], outline=colors[0], width=2)
# Car roof
draw.rectangle([cx-40, cy-40, cx+40, cy-20], outline=colors[0], width=2)
# Wheels
draw.ellipse([cx-50, cy+10, cx-30, cy+30], outline=colors[0], width=2)
draw.ellipse([cx+30, cy+10, cx+50, cy+30], outline=colors[0], width=2)
# Windows
draw.rectangle([cx-35, cy-35, cx+35, cy-25], outline=colors[1], width=1)
def _draw_house(self, draw, width, height, colors):
"""Draw a simple house sketch"""
cx, cy = width // 2, height // 2
# House base
draw.rectangle([cx-50, cy-10, cx+50, cy+40], outline=colors[0], width=2)
# Roof
draw.polygon([cx-60, cy-10, cx, cy-50, cx+60, cy-10], outline=colors[0], width=2)
# Door
draw.rectangle([cx-15, cy+10, cx+15, cy+40], outline=colors[1], width=2)
# Windows
draw.rectangle([cx-40, cy-5, cx-25, cy+10], outline=colors[1], width=1)
draw.rectangle([cx+25, cy-5, cx+40, cy+10], outline=colors[1], width=1)
def _draw_flower(self, draw, width, height, colors):
"""Draw a simple flower sketch"""
cx, cy = width // 2, height // 2
# Stem
draw.line([cx, cy+20, cx, cy+60], fill=colors[0], width=3)
# Petals
for i in range(6):
angle = i * 60 * math.pi / 180
x = cx + 25 * math.cos(angle)
y = cy + 25 * math.sin(angle)
draw.ellipse([x-8, y-8, x+8, y+8], outline=colors[0], width=2)
# Center
draw.ellipse([cx-8, cy-8, cx+8, cy+8], fill=colors[1], outline=colors[0], width=2)
# Leaves
draw.ellipse([cx-10, cy+30, cx+10, cy+50], outline=colors[0], width=2)
def _draw_tree(self, draw, width, height, colors):
"""Draw a simple tree sketch"""
cx, cy = width // 2, height // 2
# Trunk
draw.rectangle([cx-8, cy+10, cx+8, cy+60], outline=colors[0], width=2)
# Tree crown (circle)
draw.ellipse([cx-40, cy-40, cx+40, cy+20], outline=colors[0], width=2)
# Branches
for i in range(5):
angle = (i * 72 - 90) * math.pi / 180
x = cx + 30 * math.cos(angle)
y = cy + 30 * math.sin(angle)
draw.line([cx, cy, x, y], fill=colors[1], width=1)
def _draw_mountain(self, draw, width, height, colors):
"""Draw a simple mountain landscape"""
cx, cy = width // 2, height // 2
# Mountains
draw.polygon([20, cy+30, 80, cy-40, 140, cy+30], outline=colors[0], width=2)
draw.polygon([100, cy+30, 160, cy-20, 200, cy+30], outline=colors[0], width=2)
# Ground line
draw.line([0, cy+30, width, cy+30], fill=colors[1], width=1)
# Sun
draw.ellipse([width-60, 20, width-20, 60], outline=colors[1], width=2)
def _draw_abstract(self, draw, width, height, colors):
"""Draw abstract shapes for unknown prompts"""
cx, cy = width // 2, height // 2
# Random geometric shapes
for i in range(5):
x = random.randint(20, width-20)
y = random.randint(20, height-20)
size = random.randint(10, 30)
if i % 3 == 0:
draw.ellipse([x-size, y-size, x+size, y+size], outline=colors[i%len(colors)], width=2)
elif i % 3 == 1:
draw.rectangle([x-size, y-size, x+size, y+size], outline=colors[i%len(colors)], width=2)
else:
points = []
for j in range(3):
px = x + size * math.cos(j * 120 * math.pi / 180)
py = y + size * math.sin(j * 120 * math.pi / 180)
points.extend([px, py])
draw.polygon(points, outline=colors[i%len(colors)], width=2) |