import torch from PIL import Image import cv2 import numpy as np import torchvision.transforms as transforms from transformers import BlipProcessor, BlipForConditionalGeneration from PIL import Image import torch # Setup device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load BLIP model and processor blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device) def get_image_caption(image: Image.Image) -> str: image = image.convert("RGB") inputs = blip_processor(images=image, return_tensors="pt").to(device) output = blip_model.generate(**inputs) caption = blip_processor.decode(output[0], skip_special_tokens=True) return caption def generate_image_from_sketch(sketch_img: Image.Image, generator_model: torch.nn.Module, device='cpu') -> Image.Image: # Resize and convert to RGB to ensure 3 channels sketch_img = sketch_img.resize((256, 256)).convert("RGB") # Convert to tensor and normalize to [-1, 1] transform = transforms.Compose([ transforms.ToTensor(), # [0, 1] transforms.Normalize([0.5]*3, [0.5]*3) # → [-1, 1] ]) input_tensor = transform(sketch_img).unsqueeze(0).to(device) # Generate output generator_model.to(device) generator_model.eval() with torch.no_grad(): output_tensor = generator_model(input_tensor) # Denormalize output back to [0, 1] output_tensor = output_tensor.squeeze(0).cpu() * 0.5 + 0.5 output_img = transforms.ToPILImage()(output_tensor.clamp(0, 1)) return output_img def dodge_sketch(image: Image.Image) -> Image.Image: image_np = np.array(image.convert("RGB")) gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY) inverted = 255 - gray blur = cv2.GaussianBlur(inverted, (21, 21), sigmaX=0, sigmaY=0) dodge = cv2.divide(gray, 255 - blur, scale=256) edges = cv2.Laplacian(gray, cv2.CV_8U, ksize=5) edges = cv2.GaussianBlur(edges, (3, 3), 0) sketch = cv2.subtract(dodge, edges // 3) sketch = np.clip(sketch, 0, 255).astype(np.uint8) sketch = cv2.fastNlMeansDenoising(sketch, h=15, templateWindowSize=7, searchWindowSize=21) return Image.fromarray(sketch).convert("RGB") def sobel_sketch(image: Image.Image) -> Image.Image: # Convert PIL Image to NumPy array (RGB) and then to grayscale image_np = np.array(image.convert("RGB")) gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY) sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3) sobel_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3) abs_sobel_x = cv2.convertScaleAbs(sobel_x) abs_sobel_y = cv2.convertScaleAbs(sobel_y) sobel_combined = cv2.addWeighted(abs_sobel_x, 0.5, abs_sobel_y, 0.5, 0) inverted_edges = 255 - sobel_combined sketch = cv2.fastNlMeansDenoising(inverted_edges, h=15, templateWindowSize=7, searchWindowSize=21) return Image.fromarray(sketch).convert("RGB") def lattice_sketch(image: Image.Image) -> Image.Image: image_np = np.array(image.convert("RGB")) gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY) sharpen_kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]]) sharpened = cv2.filter2D(gray, -1, sharpen_kernel) inverted = 255 - sharpened blur = cv2.GaussianBlur(inverted, (21, 21), 0) dodge = cv2.divide(sharpened, 255 - blur, scale=256) sketch = cv2.equalizeHist(dodge) sketch = cv2.fastNlMeansDenoising(sketch, h=75, templateWindowSize=7, searchWindowSize=21) return Image.fromarray(sketch).convert("RGB") def canny_sketch(image: Image.Image) -> Image.Image: def gaussian_blur(image, kernel_size=5, sigma=1.4): return cv2.GaussianBlur(image, (kernel_size, kernel_size), sigma) def sobel_filters(img): Ix = cv2.Sobel(img, cv2.CV_64F, 1, 0, ksize=3) Iy = cv2.Sobel(img, cv2.CV_64F, 0, 1, ksize=3) magnitude = np.hypot(Ix, Iy) angle = np.arctan2(Iy, Ix) magnitude = magnitude / magnitude.max() * 255 return magnitude.astype(np.uint8), angle def non_max_suppression(mag, angle): M, N = mag.shape output = np.zeros((M, N), dtype=np.uint8) angle = angle * 180 / np.pi angle[angle < 0] += 180 for i in range(1, M - 1): for j in range(1, N - 1): q, r = 255, 255 a = angle[i, j] if (0 <= a < 22.5) or (157.5 <= a <= 180): q = mag[i, j + 1] r = mag[i, j - 1] elif (22.5 <= a < 67.5): q = mag[i + 1, j - 1] r = mag[i - 1, j + 1] elif (67.5 <= a < 112.5): q = mag[i + 1, j] r = mag[i - 1, j] elif (112.5 <= a < 157.5): q = mag[i - 1, j - 1] r = mag[i + 1, j + 1] output[i, j] = mag[i, j] if mag[i, j] >= q and mag[i, j] >= r else 0 return output def threshold(image, low_ratio=0.05, high_ratio=0.15): high_threshold = image.max() * high_ratio low_threshold = high_threshold * low_ratio M, N = image.shape res = np.zeros((M, N), dtype=np.uint8) strong, weak = np.uint8(255), np.uint8(75) strong_i, strong_j = np.where(image >= high_threshold) weak_i, weak_j = np.where((image <= high_threshold) & (image >= low_threshold)) res[strong_i, strong_j] = strong res[weak_i, weak_j] = weak return res, weak, strong def hysteresis(img, weak=75, strong=255): M, N = img.shape for i in range(1, M - 1): for j in range(1, N - 1): if img[i, j] == weak: if ((img[i+1, j-1] == strong) or (img[i+1, j] == strong) or (img[i+1, j+1] == strong) or (img[i, j-1] == strong) or (img[i, j+1] == strong) or (img[i-1, j-1] == strong) or (img[i-1, j] == strong) or (img[i-1, j+1] == strong)): img[i, j] = strong else: img[i, j] = 0 return img image_np = np.array(image.convert("RGB")) gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY) blurred = gaussian_blur(gray) mag, angle = sobel_filters(blurred) nms = non_max_suppression(mag, angle) thresh, weak, strong = threshold(nms) result = hysteresis(thresh, weak, strong) inverted = 255 - result return Image.fromarray(inverted).convert("RGB")