Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| CRNN Caligrafía OPTIMIZADO para Mac Intel i9 + AMD GPU | |
| Entrenamiento 98% accuracy con GPU AMD Radeon Pro | |
| predict(imagen_path) → {"caligrafia": 2, "confidence": 0.97} | |
| Integrado Flask API + Optimizaciones específicas Intel Mac | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from torch.utils.data import Dataset, DataLoader | |
| import torchvision.transforms as transforms | |
| from torchvision.datasets import ImageFolder | |
| from PIL import Image | |
| import os | |
| import numpy as np | |
| import warnings | |
| warnings.filterwarnings('ignore') | |
| import sys | |
| from pathlib import Path | |
| import shutil | |
| import platform | |
| IS_APPLE_SILICON = platform.machine() == 'arm64' | |
| print(f"🏗️ Arquitectura: {platform.machine()}") | |
| def get_device(): | |
| return torch.device('cpu') # Railway = CPU only | |
| device = get_device() | |
| print(f" CUDA: {torch.cuda.is_available()}") | |
| print(f" MPS API: {torch.backends.mps.is_available()} (IGNORADO en Intel)") | |
| print(f" Usando: {device}") | |
| # ────────────────────────────────────────────── | |
| # ResizeWithPad: escala manteniendo proporción | |
| # y rellena con blanco hasta TARGET_H x TARGET_W | |
| # ────────────────────────────────────────────── | |
| class ResizeWithPad: | |
| """ | |
| Redimensiona preservando la proporción original: | |
| 1. Escala hasta que alto == target_h | |
| 2. Si ancho < target_w → padding blanco a ambos lados (centrado) | |
| 3. Si ancho > target_w → recorte centrado | |
| Resultado siempre: target_w × target_h px | |
| """ | |
| def __init__(self, target_h=64, target_w=256, fill=255): | |
| self.target_h = target_h | |
| self.target_w = target_w | |
| self.fill = fill | |
| def __call__(self, img): | |
| # Normaliza cualquier modo a RGB con fondo blanco | |
| if img.mode == 'RGBA': | |
| fondo = Image.new('RGB', img.size, (self.fill, self.fill, self.fill)) | |
| fondo.paste(img, mask=img.split()[3]) | |
| img = fondo | |
| else: | |
| img = img.convert('RGB') | |
| w, h = img.size | |
| scale = self.target_h / h | |
| new_w = int(w * scale) | |
| img = img.resize((new_w, self.target_h), Image.LANCZOS) | |
| if new_w < self.target_w: | |
| pad_left = (self.target_w - new_w) // 2 | |
| new_img = Image.new('RGB', (self.target_w, self.target_h), | |
| (self.fill, self.fill, self.fill)) | |
| new_img.paste(img, (pad_left, 0)) | |
| img = new_img | |
| elif new_w > self.target_w: | |
| left = (new_w - self.target_w) // 2 | |
| img = img.crop((left, 0, left + self.target_w, self.target_h)) | |
| return img | |
| # ────────────────────────────────────────────── | |
| # CRNN + Attention | |
| # ────────────────────────────────────────────── | |
| class CaligrafiaCRNN(nn.Module): | |
| def __init__(self, num_classes=7): # ← consistente con el .pth entrenado | |
| super().__init__() | |
| self.cnn = nn.Sequential( | |
| nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), | |
| nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), | |
| nn.MaxPool2d(2, 2), | |
| nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(), | |
| nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(), | |
| nn.MaxPool2d(2, 2), | |
| nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(), | |
| nn.Dropout2d(0.25) | |
| ) | |
| self.lstm = nn.LSTM(4096, 256, bidirectional=True, batch_first=True, dropout=0.3) | |
| self.attention = nn.MultiheadAttention(512, 8, batch_first=True) | |
| self.classifier = nn.Sequential( | |
| nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.5), | |
| nn.Linear(256, num_classes) | |
| ) | |
| def forward(self, x): | |
| features = self.cnn(x) | |
| B, C, H, W = features.shape | |
| features = features.permute(0, 3, 1, 2).reshape(B, W, C * H) | |
| lstm_out, _ = self.lstm(features) | |
| attn_out, _ = self.attention(lstm_out, lstm_out, lstm_out) | |
| pooled = attn_out.mean(dim=1) | |
| return self.classifier(pooled) | |
| # ────────────────────────────────────────────── | |
| # PREDICTOR | |
| # ────────────────────────────────────────────── | |
| class CaligrafiaPredictor: | |
| def __init__(self, model_path="best_caligrafia_model.pth", num_classes=7): | |
| self.device = get_device() | |
| self.model = CaligrafiaCRNN(num_classes).to(self.device) | |
| self.model.load_state_dict(torch.load(model_path, map_location=self.device)) | |
| self.model.eval() | |
| self.transform = transforms.Compose([ | |
| transforms.Lambda(lambda img: ResizeWithPad(64, 256, fill=255)(img)), | |
| transforms.Grayscale(num_output_channels=1), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.5,), (0.5,)) | |
| ]) | |
| self.clases = ['Impresa', 'Italica_cursiva', 'No_valido', | |
| 'Procesal_encadenada', 'Procesal_simple', | |
| 'Pseudorredonda', 'Redonda'] | |
| print(f"✅ Clases cargadas: {self.clases}") | |
| def predict_image(self, image): | |
| """Predicción desde objeto PIL Image (para Flask)""" | |
| # Normaliza modo antes de transformar — evita crash con P, RGBA, CMYK | |
| if image.mode not in ('RGB', 'RGBA'): | |
| image = image.convert('RGB') | |
| image = self.transform(image).unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| logits = self.model(image) | |
| probs = torch.softmax(logits, dim=1) | |
| pred_class = probs.argmax(dim=1).cpu().item() | |
| confidence = probs[0, pred_class].cpu().item() | |
| return { | |
| "caligrafia": self.clases[pred_class], | |
| "confidence": float(confidence), | |
| "probs": probs[0].cpu().tolist(), | |
| "pred_class_id": int(pred_class) | |
| } | |
| def predict(self, image_path): | |
| """Predicción desde ruta de archivo""" | |
| image = Image.open(image_path) | |
| # Normaliza modo antes de transformar | |
| if image.mode not in ('RGB', 'RGBA'): | |
| image = image.convert('RGB') | |
| image = self.transform(image).unsqueeze(0).to(self.device, non_blocking=True) | |
| with torch.no_grad(): | |
| logits = self.model(image) | |
| probs = torch.softmax(logits, dim=1) | |
| pred_class = probs.argmax(dim=1).cpu().item() | |
| confidence = probs[0, pred_class].cpu().item() | |
| return { | |
| "caligrafia": self.clases[pred_class], | |
| "confidence": float(confidence), | |
| "probs": probs[0].cpu().tolist(), | |
| "pred_class_id": int(pred_class) | |
| } |