Guanc27/check_fonts

1"""2Image preprocessing for font detection.3 4Cleans up real-world screenshots and photos so they work well with the5OpenCLIP model trained on clean synthetic samples.6 7Two paths:8  - preprocess_single(): For pre-cropped regions (Chrome extension use case)9  - preprocess_for_model(): For full photos with multiple text regions (upload use case)10"""11 12import cv213import numpy as np14import torch15from PIL import Image, ImageOps16from torchvision import transforms17 18 19class FontImagePreprocessor:20    """Prepares real-world font images for the OpenCLIP embedding model."""21 22    # OpenCLIP ViT-B-32 normalization constants23    OPENCLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)24    OPENCLIP_STD = (0.26862954, 0.26130258, 0.27577711)25 26    def __init__(self, target_size: int = 224):27        self.target_size = target_size28        # Only ToTensor + Normalize — we handle resize/pad ourselves to29        # avoid OpenCLIP's default CenterCrop which cuts off wide text.30        self.tensor_transform = transforms.Compose([31            transforms.ToTensor(),32            transforms.Normalize(mean=self.OPENCLIP_MEAN, std=self.OPENCLIP_STD),33        ])34        self._ocr_reader = None  # lazy-loaded35 36    # ------------------------------------------------------------------37    # Public API38    # ------------------------------------------------------------------39 40    def preprocess_single(self, image: Image.Image, preprocess_fn=None) -> torch.Tensor:41        """Primary path for the Chrome extension.42 43        The user already selected a region, so skip text detection.44        Pipeline: deskew -> normalize background -> resize+pad -> tensor normalize.45 46        Args:47            image: PIL Image (RGB).48            preprocess_fn: Ignored (kept for interface symmetry).49 50        Returns:51            Tensor of shape ``[1, 3, 224, 224]``.52        """53        image = image.convert("RGB")54        image = self.deskew(image)55        image = self.normalize_background(image)56        image = self.resize_and_pad(image, self.target_size)57        tensor = self.tensor_transform(image)58        return tensor.unsqueeze(0)59 60    def preprocess_for_model(self, image: Image.Image, preprocess_fn=None) -> torch.Tensor:61        """Secondary path for uploaded photos with multiple text regions.62 63        Runs EasyOCR text detection, crops each detected region, then runs64        the single-region pipeline on each crop.65 66        Args:67            image: PIL Image (RGB).68            preprocess_fn: Ignored.69 70        Returns:71            Tensor of shape ``[N, 3, 224, 224]`` where N is the number of72            detected text regions (at least 1 — falls back to full image).73        """74        image = image.convert("RGB")75        boxes = self._detect_text_regions(image)76 77        if not boxes:78            # No text detected — treat the whole image as one region79            return self.preprocess_single(image)80 81        tensors = []82        for (x1, y1, x2, y2) in boxes:83            crop = image.crop((x1, y1, x2, y2))84            tensors.append(self.preprocess_single(crop))85 86        return torch.cat(tensors, dim=0)87 88    # ------------------------------------------------------------------89    # Image cleanup helpers90    # ------------------------------------------------------------------91 92    @staticmethod93    def deskew(image: Image.Image) -> Image.Image:94        """Straighten slightly rotated text using OpenCV contour analysis."""95        img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)96        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)97        _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)98 99        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)100        if not contours:101            return image102 103        # Use the largest contour to determine skew angle104        largest = max(contours, key=cv2.contourArea)105        rect = cv2.minAreaRect(largest)106        angle = rect[2]107 108        # minAreaRect returns angles in [-90, 0); normalise to [-45, 45]109        if angle < -45:110            angle += 90111        elif angle > 45:112            angle -= 90113 114        # Only correct small skews (< 15 degrees) to avoid mangling115        if abs(angle) < 0.5 or abs(angle) > 15:116            return image117 118        h, w = img_cv.shape[:2]119        center = (w // 2, h // 2)120        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)121        rotated = cv2.warpAffine(122            img_cv, rotation_matrix, (w, h),123            flags=cv2.INTER_LINEAR,124            borderMode=cv2.BORDER_REPLICATE,125        )126        return Image.fromarray(cv2.cvtColor(rotated, cv2.COLOR_BGR2RGB))127 128    @staticmethod129    def normalize_background(image: Image.Image) -> Image.Image:130        """Ensure dark-text-on-light-background and boost contrast."""131        gray = image.convert("L")132        mean_brightness = np.array(gray).mean()133 134        # If the image is predominantly dark, it's likely light text on dark bg135        if mean_brightness < 128:136            image = ImageOps.invert(image)137 138        image = ImageOps.autocontrast(image)139        return image140 141    @staticmethod142    def resize_and_pad(image: Image.Image, target: int = 224) -> Image.Image:143        """Resize longest side to *target*, paste centred on white canvas.144 145        This avoids OpenCLIP's default CenterCrop which cuts off text on146        wide or tall images.147        """148        w, h = image.size149        scale = target / max(w, h)150        new_w, new_h = int(w * scale), int(h * scale)151        image = image.resize((new_w, new_h), Image.LANCZOS)152 153        canvas = Image.new("RGB", (target, target), (255, 255, 255))154        paste_x = (target - new_w) // 2155        paste_y = (target - new_h) // 2156        canvas.paste(image, (paste_x, paste_y))157        return canvas158 159    # ------------------------------------------------------------------160    # Text detection (lazy-loaded EasyOCR)161    # ------------------------------------------------------------------162 163    def _get_ocr_reader(self):164        if self._ocr_reader is None:165            import easyocr166            self._ocr_reader = easyocr.Reader(["en"], gpu=torch.cuda.is_available())167        return self._ocr_reader168 169    def _detect_text_regions(self, image: Image.Image):170        """Return a list of (x1, y1, x2, y2) bounding boxes for text regions."""171        reader = self._get_ocr_reader()172        img_np = np.array(image)173        results = reader.readtext(img_np)174 175        boxes = []176        for (bbox, _text, _conf) in results:177            xs = [pt[0] for pt in bbox]178            ys = [pt[1] for pt in bbox]179            x1, y1 = int(min(xs)), int(min(ys))180            x2, y2 = int(max(xs)), int(max(ys))181            # Skip tiny detections182            if (x2 - x1) > 10 and (y2 - y1) > 5:183                boxes.append((x1, y1, x2, y2))184 185        return boxes186