File size: 12,164 Bytes
70410f6 |
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 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 |
from PIL import Image, ImageDraw, ImageFilter
import math
import random
import io
import base64
class EndpointHandler:
def __init__(self, path=""):
"""Initialize DiffSketchEdit handler for Hugging Face Inference API"""
pass
def __call__(self, data):
"""Edit sketch based on text prompt"""
# Extract prompt
inputs = data.get("inputs", "")
if isinstance(inputs, dict):
prompt = inputs.get("prompt", inputs.get("text", ""))
input_image_data = inputs.get("input_image", None)
edit_type = inputs.get("edit_type", "refine")
else:
prompt = str(inputs)
input_image_data = data.get("input_image", None)
edit_type = data.get("edit_type", "refine")
if not prompt:
prompt = "edit sketch"
# Handle input image if provided
input_image = None
if input_image_data:
try:
if isinstance(input_image_data, str):
# Base64 encoded image
img_data = base64.b64decode(input_image_data)
input_image = Image.open(io.BytesIO(img_data)).convert('RGB')
input_image = input_image.resize((224, 224))
except Exception:
input_image = None
# Generate/edit sketch
image = self.edit_sketch(prompt, input_image, edit_type)
# Return PIL Image directly for HF Inference API
return image
def edit_sketch(self, prompt, input_image=None, edit_type="refine"):
"""Edit or create a sketch based on the prompt"""
if input_image is None:
# Create initial sketch
img = self._create_initial_sketch(prompt)
else:
img = input_image.copy()
# Apply editing based on type
if edit_type == "refine":
img = self._refine_sketch(img, prompt)
elif edit_type == "style_transfer":
img = self._apply_style_transfer(img, prompt)
elif edit_type == "add_details":
img = self._add_details(img, prompt)
elif edit_type == "color":
img = self._apply_coloring(img, prompt)
else:
# Default refinement
img = self._refine_sketch(img, prompt)
return img
def _create_initial_sketch(self, prompt):
"""Create initial sketch based on prompt"""
img = Image.new('RGB', (224, 224), 'white')
draw = ImageDraw.Draw(img)
prompt_lower = prompt.lower()
if any(word in prompt_lower for word in ['house', 'building', 'home']):
self._draw_house_sketch(draw)
elif any(word in prompt_lower for word in ['tree', 'plant', 'nature']):
self._draw_tree_sketch(draw)
elif any(word in prompt_lower for word in ['face', 'portrait', 'person']):
self._draw_face_sketch(draw)
elif any(word in prompt_lower for word in ['car', 'vehicle']):
self._draw_car_sketch(draw)
elif any(word in prompt_lower for word in ['flower', 'bloom']):
self._draw_flower_sketch(draw)
else:
self._draw_abstract_sketch(draw, prompt)
return img
def _draw_house_sketch(self, draw):
"""Draw basic house sketch"""
# Base
draw.rectangle([50, 120, 174, 180], outline='black', width=2)
# Roof
draw.polygon([(50, 120), (112, 80), (174, 120)], outline='black', width=2)
# Door
draw.rectangle([100, 150, 124, 180], outline='black', width=2)
# Windows
draw.rectangle([70, 140, 90, 160], outline='black', width=2)
draw.rectangle([134, 140, 154, 160], outline='black', width=2)
def _draw_tree_sketch(self, draw):
"""Draw basic tree sketch"""
# Trunk
draw.rectangle([105, 140, 119, 200], fill='brown', outline='black', width=2)
# Crown
draw.ellipse([70, 80, 154, 150], outline='green', width=3)
# Branches
for angle in range(0, 360, 60):
x = 112 + 25 * math.cos(math.radians(angle))
y = 115 + 25 * math.sin(math.radians(angle))
draw.line([112, 115, x, y], fill='brown', width=2)
def _draw_face_sketch(self, draw):
"""Draw basic face sketch"""
center = (112, 112)
# Head outline
draw.ellipse([center[0]-40, center[1]-50, center[0]+40, center[1]+30],
outline='black', width=2)
# Eyes
draw.ellipse([center[0]-20, center[1]-20, center[0]-10, center[1]-10],
outline='black', width=2)
draw.ellipse([center[0]+10, center[1]-20, center[0]+20, center[1]-10],
outline='black', width=2)
# Nose
draw.line([center[0], center[1]-5, center[0]-3, center[1]+5], fill='black', width=2)
# Mouth
draw.arc([center[0]-15, center[1]+5, center[0]+15, center[1]+20], 0, 180, fill='black', width=2)
def _draw_car_sketch(self, draw):
"""Draw basic car sketch"""
# Body
draw.rectangle([50, 120, 174, 160], outline='black', width=2)
# Roof
draw.rectangle([70, 100, 154, 120], outline='black', width=2)
# Wheels
draw.ellipse([60, 150, 80, 170], outline='black', width=2)
draw.ellipse([144, 150, 164, 170], outline='black', width=2)
# Windows
draw.rectangle([75, 105, 100, 115], outline='black', width=1)
draw.rectangle([124, 105, 149, 115], outline='black', width=1)
def _draw_flower_sketch(self, draw):
"""Draw basic flower sketch"""
center = (112, 112)
# Stem
draw.line([center[0], center[1]+20, center[0], 200], fill='green', width=4)
# Petals
for angle in range(0, 360, 45):
x = center[0] + 25 * math.cos(math.radians(angle))
y = center[1] + 25 * math.sin(math.radians(angle))
draw.ellipse([x-8, y-15, x+8, y+5], outline='black', width=2)
# Center
draw.ellipse([center[0]-8, center[1]-8, center[0]+8, center[1]+8],
outline='black', width=2)
def _draw_abstract_sketch(self, draw, prompt):
"""Draw abstract sketch"""
prompt_hash = hash(prompt) % 100
for i in range(5):
x1 = (i * 40 + prompt_hash) % 180 + 22
y1 = (i * 30 + prompt_hash) % 160 + 32
x2 = x1 + 40 + (i * 10) % 30
y2 = y1 + 60 + (i * 15) % 40
draw.ellipse([x1, y1, x2, y2], outline='black', width=2)
def _refine_sketch(self, img, prompt):
"""Refine the sketch with enhanced details"""
# Apply sharpening
img = img.filter(ImageFilter.SHARPEN)
draw = ImageDraw.Draw(img)
# Add refinement details based on prompt
prompt_lower = prompt.lower()
if "house" in prompt_lower:
# Add roof tiles
for y in range(80, 120, 5):
for x in range(50 + (y-80)//2, 174 - (y-80)//2, 10):
draw.line([x, y, x+5, y], fill='gray', width=1)
# Add door handle
draw.ellipse([120, 164, 122, 166], fill='black')
elif "tree" in prompt_lower:
# Add leaf details
for i in range(10):
x = 70 + random.randint(0, 84)
y = 80 + random.randint(0, 70)
draw.ellipse([x-2, y-2, x+2, y+2], fill='green')
elif "face" in prompt_lower:
# Add facial details
center = (112, 112)
# Eyebrows
draw.arc([center[0]-25, center[1]-35, center[0]-5, center[1]-25], 0, 180, fill='black', width=2)
draw.arc([center[0]+5, center[1]-35, center[0]+25, center[1]-25], 0, 180, fill='black', width=2)
# Pupils
draw.ellipse([center[0]-17, center[1]-17, center[0]-13, center[1]-13], fill='black')
draw.ellipse([center[0]+13, center[1]-17, center[0]+17, center[1]-13], fill='black')
return img
def _apply_style_transfer(self, img, prompt):
"""Apply style transfer effects"""
prompt_lower = prompt.lower()
if "cartoon" in prompt_lower:
# Cartoon style - enhance colors and add outlines
img = img.filter(ImageFilter.EDGE_ENHANCE_MORE)
# Increase saturation
from PIL import ImageEnhance
enhancer = ImageEnhance.Color(img)
img = enhancer.enhance(1.5)
elif "realistic" in prompt_lower:
# Realistic style - add shading and texture
img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
# Add subtle noise for texture
import numpy as np
pixels = np.array(img)
noise = np.random.normal(0, 3, pixels.shape)
pixels = np.clip(pixels + noise, 0, 255).astype(np.uint8)
img = Image.fromarray(pixels)
elif "watercolor" in prompt_lower:
# Watercolor style - soft edges and bleeding
img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
from PIL import ImageEnhance
enhancer = ImageEnhance.Color(img)
img = enhancer.enhance(0.8)
return img
def _add_details(self, img, prompt):
"""Add contextual details to the sketch"""
draw = ImageDraw.Draw(img)
prompt_lower = prompt.lower()
if "landscape" in prompt_lower:
# Add clouds
for i in range(3):
x = 40 + i * 60
y = 30 + i * 10
draw.ellipse([x, y, x+30, y+15], fill='lightgray', outline='gray')
# Add birds
for i in range(2):
x = 60 + i * 80
y = 50 + i * 5
draw.arc([x, y, x+10, y+5], 0, 180, fill='black', width=1)
elif "portrait" in prompt_lower:
# Add hair details
center = (112, 112)
for i in range(15):
x = center[0] + random.randint(-50, 50)
y = center[1] - 60 + random.randint(0, 30)
draw.line([x, y, x + random.randint(-5, 5), y + random.randint(10, 30)],
fill='brown', width=1)
elif "building" in prompt_lower:
# Add architectural details
# Window frames
draw.rectangle([69, 139, 91, 161], outline='black', width=1)
draw.rectangle([133, 139, 155, 161], outline='black', width=1)
# Chimney
draw.rectangle([140, 70, 150, 90], outline='black', width=2)
return img
def _apply_coloring(self, img, prompt):
"""Apply coloring to the sketch"""
# Convert to RGBA for transparency effects
img = img.convert('RGBA')
# Create color overlay
color_overlay = Image.new('RGBA', img.size, (0, 0, 0, 0))
draw = ImageDraw.Draw(color_overlay)
prompt_lower = prompt.lower()
if "sunset" in prompt_lower:
# Sunset colors
draw.rectangle([0, 0, 224, 112], fill=(255, 200, 100, 50))
draw.rectangle([0, 112, 224, 224], fill=(255, 150, 50, 30))
elif "nature" in prompt_lower:
# Nature colors
draw.rectangle([0, 140, 224, 224], fill=(100, 200, 100, 40)) # Green ground
draw.rectangle([0, 0, 224, 140], fill=(150, 200, 255, 30)) # Blue sky
elif "warm" in prompt_lower:
# Warm color palette
draw.rectangle([0, 0, 224, 224], fill=(255, 200, 150, 25))
elif "cool" in prompt_lower:
# Cool color palette
draw.rectangle([0, 0, 224, 224], fill=(150, 200, 255, 25))
# Blend with original
img = Image.alpha_composite(img, color_overlay)
img = img.convert('RGB')
return img |