jree423
/

diffsketchedit

@@ -1,307 +0,0 @@
-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