| import io |
| import copy |
| import hashlib |
| import hmac |
| import json |
| import os |
| import random |
| import re |
| import shutil |
| import tempfile |
| import threading |
| import uuid |
| import zipfile |
| from collections import Counter |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| import torch |
| from fastapi import FastAPI, File, Form, Header, HTTPException, UploadFile |
| from huggingface_hub import HfApi |
| from PIL import Image, ImageOps |
| from torch import nn |
| from torch.utils.data import DataLoader, Dataset |
| from torchvision.models import efficientnet_b0 |
| from torchvision.transforms import v2 as transforms |
|
|
| app = FastAPI(title="Embroidery Complexity API", version="1.0.0") |
| MODEL_PATH = Path(__file__).parent / "complexidade_best.pth" |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| checkpoint = torch.load(MODEL_PATH, map_location=device) |
| classes = checkpoint["classes"] |
| image_size = checkpoint.get("image_size", 224) |
|
|
|
|
| def build_model(model_checkpoint: dict, allow_classifier_mismatch: bool = False) -> torch.nn.Module: |
| model_classes = model_checkpoint["classes"] |
| trained_model = efficientnet_b0(weights=None) |
| trained_model.classifier[1] = torch.nn.Linear(trained_model.classifier[1].in_features, len(model_classes)) |
| state = model_checkpoint["model_state"] |
| if allow_classifier_mismatch: |
| current_state = trained_model.state_dict() |
| state = {key: value for key, value in state.items() if key in current_state and current_state[key].shape == value.shape} |
| trained_model.load_state_dict(state, strict=False) |
| else: |
| trained_model.load_state_dict(state) |
| return trained_model.to(device).eval() |
|
|
|
|
| model = build_model(checkpoint) |
|
|
|
|
| def make_transform(size: int): |
| return transforms.Compose([ |
| transforms.Resize((size, size)), |
| transforms.ToImage(), |
| transforms.ToDtype(torch.float32, scale=True), |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
| ]) |
|
|
|
|
| transform = make_transform(image_size) |
| MODEL_LOCK = threading.RLock() |
| TRAIN_EXECUTOR = ThreadPoolExecutor(max_workers=1) |
| TRAINING_JOBS: dict[str, dict] = {} |
| MAX_DATASET_BYTES = 500 * 1024 * 1024 |
| MAX_EXTRACTED_BYTES = 2 * 1024 * 1024 * 1024 |
| MAX_DATASET_FILES = 5000 |
|
|
|
|
| def extract_palette(image: Image.Image, quantity: int = 12) -> list[dict]: |
| small = image.copy() |
| small.thumbnail((512, 512), Image.Resampling.BILINEAR) |
| quantized = small.quantize(colors=quantity, method=Image.Quantize.MEDIANCUT) |
| palette = quantized.getpalette() or [] |
| colors = quantized.getcolors(maxcolors=quantity * 512 * 512) or [] |
| total = sum(count for count, _ in colors) or 1 |
| result = [] |
| for count, index in sorted(colors, reverse=True): |
| percentage = count / total |
| if percentage < 0.01: |
| continue |
| base = index * 3 |
| rgb = tuple(int(palette[base + offset]) for offset in range(3)) |
| result.append({"hex": "#%02x%02x%02x" % rgb, "rgb": list(rgb), "percentage": round(percentage, 4)}) |
| if len(result) >= 8: |
| break |
| return result |
|
|
|
|
| def _training_token_is_valid(provided_token: str | None) -> bool: |
| expected = os.getenv("TRAINING_API_TOKEN", "").strip() |
| return bool(expected and provided_token and hmac.compare_digest(provided_token, expected)) |
|
|
|
|
| def _safe_extract_dataset(zip_bytes: bytes, destination: Path) -> Path: |
| """Extract only a bounded ZIP and return the folder containing the labels.""" |
| if len(zip_bytes) > MAX_DATASET_BYTES: |
| raise ValueError("Dataset ZIP is larger than the 500 MB limit.") |
| destination.mkdir(parents=True, exist_ok=True) |
| with zipfile.ZipFile(io.BytesIO(zip_bytes)) as archive: |
| entries = [item for item in archive.infolist() if not item.is_dir()] |
| if len(entries) > MAX_DATASET_FILES: |
| raise ValueError("Dataset contains too many files.") |
| total_size = sum(max(0, item.file_size) for item in entries) |
| if total_size > MAX_EXTRACTED_BYTES: |
| raise ValueError("Uncompressed dataset is larger than the 2 GB limit.") |
| destination_root = destination.resolve() |
| for item in entries: |
| target = (destination / item.filename).resolve() |
| if destination_root not in target.parents: |
| raise ValueError("Dataset contains an unsafe ZIP path.") |
| target.parent.mkdir(parents=True, exist_ok=True) |
| with archive.open(item) as source, target.open("wb") as output: |
| shutil.copyfileobj(source, output, length=1024 * 1024) |
|
|
| candidates = [destination, destination / "dataset" / "rotuladas", destination / "rotuladas"] |
| for candidate in candidates: |
| if all((candidate / label).is_dir() for label in ("baixo_risco", "medio_risco", "alto_risco")): |
| return candidate |
| raise ValueError("ZIP must contain baixo_risco/, medio_risco/ and alto_risco/ folders.") |
|
|
|
|
| def _order_id(path: Path) -> str: |
| match = re.match(r"(\d+)_", path.name) |
| return match.group(1) if match else path.stem |
|
|
|
|
| class TrainingImageDataset(Dataset): |
| def __init__(self, samples: list[tuple[Path, int]], image_transform): |
| self.samples = samples |
| self.transform = image_transform |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
| def __getitem__(self, index): |
| path, label = self.samples[index] |
| previous_limit = Image.MAX_IMAGE_PIXELS |
| Image.MAX_IMAGE_PIXELS = None |
| try: |
| with Image.open(path) as source: |
| image = ImageOps.exif_transpose(source).convert("RGB") |
| finally: |
| Image.MAX_IMAGE_PIXELS = previous_limit |
| return self.transform(image), label |
|
|
|
|
| def _collect_training_samples(root: Path) -> tuple[list[tuple[Path, int]], list[str]]: |
| class_names = ["baixo_risco", "medio_risco", "alto_risco"] |
| extensions = {".png", ".jpg", ".jpeg", ".webp", ".gif"} |
| samples: list[tuple[Path, int]] = [] |
| for label, class_name in enumerate(class_names): |
| folder = root / class_name |
| for path in sorted(folder.rglob("*")): |
| if path.is_file() and path.suffix.lower() in extensions: |
| samples.append((path, label)) |
| counts = Counter(label for _, label in samples) |
| if any(counts[label] < 6 for label in range(len(class_names))): |
| raise ValueError("Each class needs at least 6 images.") |
| groups: dict[str, list[tuple[Path, int]]] = {} |
| for sample in samples: |
| groups.setdefault(_order_id(sample[0]), []).append(sample) |
| if len(groups) < 9: |
| raise ValueError("Dataset needs at least 9 different orders, with 3 per class.") |
| for group in groups.values(): |
| if len({label for _, label in group}) != 1: |
| raise ValueError(f"The same order has different labels: {_order_id(group[0][0])}.") |
| return samples, class_names |
|
|
|
|
| def _split_training_samples(samples: list[tuple[Path, int]]) -> tuple[list, list]: |
| groups: dict[str, list[tuple[Path, int]]] = {} |
| for sample in samples: |
| groups.setdefault(_order_id(sample[0]), []).append(sample) |
| by_label: dict[int, list[list[tuple[Path, int]]]] = {} |
| for group in groups.values(): |
| by_label.setdefault(group[0][1], []).append(group) |
| rng = random.Random(42) |
| training_groups = [] |
| validation_groups = [] |
| for label_groups in by_label.values(): |
| rng.shuffle(label_groups) |
| validation_groups.append(label_groups[0]) |
| training_groups.extend(label_groups[1:]) |
| flatten = lambda selected: [sample for group in selected for sample in group] |
| return flatten(training_groups), flatten(validation_groups) |
|
|
|
|
| def _evaluate_training(model_to_test, loader, criterion, class_count: int) -> tuple[float, float, float]: |
| model_to_test.eval() |
| total_loss = 0.0 |
| targets: list[int] = [] |
| predictions: list[int] = [] |
| with torch.inference_mode(): |
| for images, labels in loader: |
| images, labels = images.to(device), labels.to(device) |
| logits = model_to_test(images) |
| total_loss += criterion(logits, labels).item() * labels.size(0) |
| predictions.extend(logits.argmax(1).cpu().tolist()) |
| targets.extend(labels.cpu().tolist()) |
| accuracy = sum(actual == predicted for actual, predicted in zip(targets, predictions)) / max(1, len(targets)) |
| f1_values = [] |
| for label in range(class_count): |
| true_positive = sum(actual == label and predicted == label for actual, predicted in zip(targets, predictions)) |
| false_positive = sum(actual != label and predicted == label for actual, predicted in zip(targets, predictions)) |
| false_negative = sum(actual == label and predicted != label for actual, predicted in zip(targets, predictions)) |
| precision = true_positive / max(1, true_positive + false_positive) |
| recall = true_positive / max(1, true_positive + false_negative) |
| f1_values.append(2 * precision * recall / max(1e-9, precision + recall)) |
| return total_loss / max(1, len(loader.dataset)), sum(f1_values) / len(f1_values), accuracy |
|
|
|
|
| def _train_and_publish(job_id: str, zip_bytes: bytes, epochs: int, learning_rate: float) -> None: |
| global model, checkpoint, classes, image_size, transform |
| temporary_root = Path(tempfile.mkdtemp(prefix=f"training_{job_id}_")) |
| try: |
| TRAINING_JOBS[job_id].update({"status": "running", "message": "Extracting dataset."}) |
| dataset_root = _safe_extract_dataset(zip_bytes, temporary_root / "dataset") |
| samples, class_names = _collect_training_samples(dataset_root) |
| training_samples, validation_samples = _split_training_samples(samples) |
| if not training_samples or not validation_samples: |
| raise ValueError("The dataset split produced an empty training or validation set.") |
|
|
| weights = checkpoint.get("image_size", 224) |
| image_transform = make_transform(int(weights)) |
| train_loader = DataLoader(TrainingImageDataset(training_samples, image_transform), batch_size=16, shuffle=True, num_workers=0) |
| validation_loader = DataLoader(TrainingImageDataset(validation_samples, image_transform), batch_size=16, shuffle=False, num_workers=0) |
|
|
| candidate = build_model({"classes": class_names, "model_state": checkpoint["model_state"]}, allow_classifier_mismatch=True) |
| counts = Counter(label for _, label in training_samples) |
| class_weights = torch.tensor( |
| [len(training_samples) / (len(class_names) * max(1, counts[label])) for label in range(len(class_names))], |
| dtype=torch.float32, |
| device=device, |
| ) |
| criterion = nn.CrossEntropyLoss(weight=class_weights) |
| optimizer = torch.optim.AdamW(candidate.parameters(), lr=learning_rate, weight_decay=1e-4) |
| best_loss = float("inf") |
| best_state = None |
| no_improvement = 0 |
|
|
| for epoch in range(1, epochs + 1): |
| candidate.train() |
| running_loss = 0.0 |
| for images, labels in train_loader: |
| images, labels = images.to(device), labels.to(device) |
| optimizer.zero_grad() |
| loss = criterion(candidate(images), labels) |
| loss.backward() |
| optimizer.step() |
| running_loss += loss.item() * labels.size(0) |
| validation_loss, validation_f1, validation_accuracy = _evaluate_training(candidate, validation_loader, criterion, len(class_names)) |
| TRAINING_JOBS[job_id].update({ |
| "epoch": epoch, |
| "epochs": epochs, |
| "train_loss": round(running_loss / max(1, len(train_loader.dataset)), 5), |
| "validation_loss": round(validation_loss, 5), |
| "validation_f1": round(validation_f1, 5), |
| "validation_accuracy": round(validation_accuracy, 5), |
| }) |
| if validation_loss < best_loss - 0.001: |
| best_loss = validation_loss |
| best_state = copy.deepcopy(candidate.state_dict()) |
| no_improvement = 0 |
| else: |
| no_improvement += 1 |
| if no_improvement >= 5: |
| break |
|
|
| if best_state is None: |
| raise RuntimeError("Training did not produce a valid model.") |
| candidate.load_state_dict(best_state) |
| new_checkpoint = {"model_state": candidate.state_dict(), "classes": class_names, "image_size": int(weights)} |
| local_model_path = temporary_root / "complexidade_best.pth" |
| torch.save(new_checkpoint, local_model_path) |
|
|
| hf_token = os.getenv("HF_TOKEN", "").strip() or os.getenv("HUGGINGFACE_HUB_TOKEN", "").strip() |
| repo_id = os.getenv("SPACE_ID", "vitordigitizing/embroidery-complexity-api").strip() |
| if not hf_token: |
| raise RuntimeError("HF_TOKEN Secret is not configured with write permission.") |
| TRAINING_JOBS[job_id]["message"] = "Uploading the new model to Hugging Face." |
| HfApi(token=hf_token).upload_file( |
| path_or_fileobj=str(local_model_path), |
| path_in_repo="complexidade_best.pth", |
| repo_id=repo_id, |
| repo_type="space", |
| commit_message=f"Update complexity model from training job {job_id}", |
| ) |
| with MODEL_LOCK: |
| checkpoint = new_checkpoint |
| classes = class_names |
| image_size = int(weights) |
| transform = make_transform(image_size) |
| model = candidate.to(device).eval() |
| TRAINING_JOBS[job_id].update({"status": "completed", "message": "Model trained and uploaded to Hugging Face."}) |
| except Exception as error: |
| TRAINING_JOBS[job_id].update({"status": "failed", "message": str(error)}) |
| finally: |
| shutil.rmtree(temporary_root, ignore_errors=True) |
|
|
|
|
| @app.get("/") |
| def root(): |
| return { |
| "status": "ok", |
| "service": "Embroidery Complexity API", |
| "model": "embroidery-complexity", |
| "endpoints": { |
| "health": "/health", |
| "predict": "POST /predict with multipart field 'image'", |
| "train": "POST /train with a protected ZIP dataset", |
| "training_status": "GET /train/{job_id}", |
| }, |
| } |
|
|
|
|
| @app.get("/health") |
| def health(): |
| return {"status": "ok", "model": "embroidery-complexity", "device": str(device)} |
|
|
|
|
| @app.post("/train", status_code=202) |
| async def train( |
| dataset: UploadFile = File(...), |
| x_training_token: str | None = Header(default=None), |
| epochs: int = Form(default=20), |
| learning_rate: float = Form(default=0.0001), |
| ): |
| """Start a protected fine-tuning job and publish the best checkpoint.""" |
| if not _training_token_is_valid(x_training_token): |
| raise HTTPException(status_code=401, detail="Invalid training token.") |
| if not os.getenv("HF_TOKEN", "").strip() and not os.getenv("HUGGINGFACE_HUB_TOKEN", "").strip(): |
| raise HTTPException(status_code=503, detail="HF_TOKEN Secret is not configured.") |
| if epochs < 1 or epochs > 50: |
| raise HTTPException(status_code=400, detail="epochs must be between 1 and 50.") |
| if learning_rate <= 0 or learning_rate > 0.01: |
| raise HTTPException(status_code=400, detail="learning_rate must be between 0 and 0.01.") |
| if not (dataset.filename or "").lower().endswith(".zip"): |
| raise HTTPException(status_code=400, detail="Send the dataset as a .zip file.") |
| zip_bytes = await dataset.read(MAX_DATASET_BYTES + 1) |
| if len(zip_bytes) > MAX_DATASET_BYTES: |
| raise HTTPException(status_code=413, detail="Dataset ZIP is larger than the 500 MB limit.") |
| job_id = uuid.uuid4().hex |
| TRAINING_JOBS[job_id] = {"status": "queued", "message": "Training job queued."} |
| TRAIN_EXECUTOR.submit(_train_and_publish, job_id, zip_bytes, epochs, learning_rate) |
| return {"status": "queued", "job_id": job_id, "status_url": f"/train/{job_id}"} |
|
|
|
|
| @app.get("/train/{job_id}") |
| def training_status(job_id: str, x_training_token: str | None = Header(default=None)): |
| if not _training_token_is_valid(x_training_token): |
| raise HTTPException(status_code=401, detail="Invalid training token.") |
| job = TRAINING_JOBS.get(job_id) |
| if job is None: |
| raise HTTPException(status_code=404, detail="Training job not found.") |
| return {"job_id": job_id, **job} |
|
|
|
|
| @app.post("/predict") |
| async def predict(image: UploadFile = File(...)): |
| if image.content_type not in {"image/png", "image/jpeg", "image/webp"}: |
| raise HTTPException(status_code=400, detail="Only PNG, JPG and WEBP images are supported.") |
| try: |
| content = await image.read() |
| original = ImageOps.exif_transpose(Image.open(io.BytesIO(content))).convert("RGB") |
| tensor = transform(original).unsqueeze(0).to(device) |
| with torch.inference_mode(): |
| probabilities = torch.softmax(model(tensor), dim=1)[0] |
| index = int(probabilities.argmax()) |
| return { |
| "status": "success", |
| "result": classes[index], |
| "confidence": round(float(probabilities[index]) * 100, 2), |
| "probabilities": {name: round(float(probabilities[pos]) * 100, 2) for pos, name in enumerate(classes)}, |
| "colors_palette": extract_palette(original), |
| } |
| except Exception as error: |
| raise HTTPException(status_code=400, detail=f"Could not analyze image: {error}") from error |
|
|
