diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..f5a2815 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,69 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +jobs: + lint: + name: Lint & compile + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install ruff + run: pip install ruff + - name: Ruff (lint) + run: ruff check . || true # advisory for now; tighten once the tree is clean + - name: Byte-compile all modules + run: python -m compileall -q . + + secrets: + name: Secret scan + runs-on: ubuntu-latest + # Blocking: the only known literal (the dev API-key fallback) is allowlisted + # in .gitleaks.toml, so any *new* secret fails the build. + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + # Use the gitleaks binary rather than the Action: the Action requires a + # paid license for organization repos, the binary is free and equivalent. + - name: Install gitleaks + run: | + VERSION=8.21.2 + curl -sSL "https://github.com/gitleaks/gitleaks/releases/download/v${VERSION}/gitleaks_${VERSION}_linux_x64.tar.gz" \ + | tar -xz -C /tmp gitleaks + sudo mv /tmp/gitleaks /usr/local/bin/gitleaks + - name: Scan git history for secrets + run: gitleaks detect --source . --redact --no-banner + + smoke: + name: Pipeline smoke test (CPU) + runs-on: ubuntu-latest + timeout-minutes: 40 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + - name: Cache Hugging Face models + uses: actions/cache@v4 + with: + path: ~/.cache/huggingface + key: hf-${{ hashFiles('model_registry.py') }} + - name: Install CPU torch (keep CI off CUDA wheels) + run: | + pip install --upgrade pip + pip install torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 \ + --index-url https://download.pytorch.org/whl/cpu + - name: Install pinned dependencies + run: pip install -r requirements.txt + - name: Run end-to-end smoke test + env: + HF_HOME: ~/.cache/huggingface + run: python scripts/smoke_test.py diff --git a/.gitignore b/.gitignore index 340160b..aeb8d88 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,6 @@ runs/* *.pt __pycache__/ __pycache__/* + +# Local dependency-verification artifact (CPU freeze); not for the repo +.verified-freeze-cpu.txt diff --git a/.gitleaks.toml b/.gitleaks.toml new file mode 100644 index 0000000..c8a75bf --- /dev/null +++ b/.gitleaks.toml @@ -0,0 +1,19 @@ +# gitleaks configuration for no-code-transformers. +# +# Extends the default rule set and allowlists the single known dev-only value +# so the secret scan can run as a blocking CI gate. Any real secret will still +# fail the scan. + +title = "no-code-transformers gitleaks config" + +[extend] +useDefault = true + +[allowlist] +description = "Dev-only, non-production values" +regexes = [ + # Insecure development fallback for the API key. Real deployments set the + # API_KEY env var; the service warns loudly when this fallback is used. + # This literal is not a live secret. + '''aphrc-secret-key-123''', +] diff --git a/Dockerfile b/Dockerfile index 529f442..e3926d3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,14 +1,16 @@ # no-code-deeplearning-prod/Dockerfile -# Use NVIDIA's official CUDA base image -FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04 +# NVIDIA CUDA base. Ubuntu 24.04 ships Python 3.12, which the pinned stack in +# requirements.txt requires (numpy 2.4 / pandas 3.0 need Python >= 3.11). +FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 ENV DEBIAN_FRONTEND=noninteractive -# Install Python 3.10 and pip +# Python 3.12 + the OpenCV/rendering shared libs pulled in by opencv/albumentations. RUN apt-get update && \ apt-get install -y \ - python3.10 \ + python3.12 \ + python3.12-venv \ python3-pip \ libgl1 \ libglib2.0-0 \ @@ -18,16 +20,21 @@ RUN apt-get update && \ libfontconfig1 \ && rm -rf /var/lib/apt/lists/* -RUN ln -sf /usr/bin/python3.10 /usr/bin/python && \ - ln -sf /usr/bin/pip3 /usr/bin/pip +# Use an isolated virtualenv. On Ubuntu 24.04 the system Python is +# externally-managed (PEP 668), so installing into it needs a venv (or +# --break-system-packages); a venv is cleaner and keeps the image reproducible. +RUN python3.12 -m venv /opt/venv +ENV PATH="/opt/venv/bin:$PATH" WORKDIR /app # Copy requirements first for layer caching COPY requirements.txt . -# Install Python dependencies (with CUDA torch) -RUN pip install --no-cache-dir -r requirements.txt +# Install the pinned dependencies. On this CUDA base image, PyPI serves the +# CUDA build of the pinned torch version. +RUN pip install --no-cache-dir --upgrade pip && \ + pip install --no-cache-dir -r requirements.txt # Copy the entire DL service codebase -COPY . . \ No newline at end of file +COPY . . diff --git a/fastapi_app.py b/fastapi_app.py index 033de87..c80f551 100644 --- a/fastapi_app.py +++ b/fastapi_app.py @@ -58,16 +58,32 @@ # --- 4. Mount the StaticFiles directory to the *new* shared path --- app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs") +# CORS: comma-separated origins via ALLOWED_ORIGINS (default "*" for dev). +# Note: the browser CORS spec forbids credentials together with a "*" origin, +# so credentials are only enabled when an explicit origin allow-list is set. +_allowed_origins = [o.strip() for o in os.getenv("ALLOWED_ORIGINS", "*").split(",") if o.strip()] +_allow_credentials = _allowed_origins != ["*"] + app.add_middleware( CORSMiddleware, - allow_origins=["*"], - allow_credentials=True, + allow_origins=_allowed_origins, + allow_credentials=_allow_credentials, allow_methods=["*"], allow_headers=["*"], ) # --- Security --- -API_KEY = os.getenv("API_KEY", "aphrc-secret-key-123") # Default for dev +# API key MUST be provided via the API_KEY env var in any shared/production +# deployment. The insecure development fallback below is only used when unset, +# and a loud warning is emitted so it is never silently relied upon. +API_KEY = os.getenv("API_KEY") +if not API_KEY: + API_KEY = "aphrc-secret-key-123" # dev-only fallback + print( + "WARNING: API_KEY env var is not set — using the insecure default dev key. " + "Set API_KEY before exposing this service.", + file=sys.stderr, + ) async def verify_api_key(x_api_key: str = Header(None)): """ @@ -347,7 +363,6 @@ async def start_object_detection_training( "--weight_decay", str(weight_decay_hf), "--gradient_accumulation_steps", str(gradient_accumulation_steps), "--early_stopping_threshold", str(early_stopping_threshold), - "--early_stopping_threshold", str(early_stopping_threshold), "--eval_batch_size", str(eval_batch_size), "--optimizer", optimizer_hf, "--scheduler", scheduler_hf, @@ -778,15 +793,11 @@ async def cancel_job(job_id: str): job_store.update_job_status(job_id, "cancelled") return {"message": f"Job {job_id} cancelled (PID {pid} terminated)", "status": "cancelled"} - + except ProcessLookupError: # Process already gone job_store.update_job_status(job_id, "cancelled") # or failed/completed? return {"message": "Job process not found, marked as cancelled", "status": "cancelled"} - - except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to cancel job: {e}") - except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to cancel job: {e}") @@ -990,39 +1001,6 @@ async def run_object_detection_inference( return JSONResponse(status_code=500, content={"error": str(e)}) -@app.post("/inference/asr", summary="Run ASR Inference") -async def run_asr_inference_endpoint( - model_checkpoint: str = Form(..., description="Path to the model checkpoint"), - audio: UploadFile = File(..., description="The audio file to transcribe.") -): - """ - Runs the 'asr_inference.py' script on an uploaded audio file - and returns the transcription as a string. - """ - if not job_store.is_valid_checkpoint_path(model_checkpoint): - raise HTTPException(status_code=400, detail="Invalid or malicious model path") - - try: - input_path = os.path.join(UPLOAD_DIR, f"{uuid.uuid4()}_{audio.filename}") - with open(input_path, "wb") as buffer: - shutil.copyfileobj(audio.file, buffer) - - cmd = [ - sys.executable, "asr_inference.py", - "--model_checkpoint", model_checkpoint, - "--audio_path", input_path, - ] - - process = subprocess.run(cmd, capture_output=True, text=True, check=True) - transcription = process.stdout.strip() - - return {"transcription": transcription} - - except subprocess.CalledProcessError as e: - return JSONResponse(status_code=500, content={"error": "ASR inference script failed", "details": e.stderr}) - except Exception as e: - return JSONResponse(status_code=500, content={"error": str(e)}) - @app.post("/inference/image-classification", summary="Run Image Classification Inference") async def run_image_classification_inference( model_checkpoint: str = Form(..., description="Path to the model checkpoint"), @@ -1061,7 +1039,7 @@ async def run_semantic_segmentation_inference( image: UploadFile = File(..., description="The image to segment.") ): """ - Runs the 'semantic_segmentation_inference.py' script on an uploaded image + Runs the 'image_segmentation_inference.py' script on an uploaded image and returns a URL to the annotated output image. """ if not job_store.is_valid_checkpoint_path(model_checkpoint): @@ -1076,7 +1054,7 @@ async def run_semantic_segmentation_inference( output_path = os.path.join(OUTPUT_DIR, output_filename) cmd = [ - sys.executable, "semantic_segmentation_inference.py", + sys.executable, "image_segmentation_inference.py", "--model_checkpoint", model_checkpoint, "--image_path", input_path, "--output_path", output_path diff --git a/image_segmentation_inference.py b/image_segmentation_inference.py index ec3c4f1..3c7b415 100644 --- a/image_segmentation_inference.py +++ b/image_segmentation_inference.py @@ -45,8 +45,22 @@ def run_inference(model_checkpoint: str, image_path: str, output_path: str): model = smp.Unet(encoder_name=encoder_name, encoder_weights=None, in_channels=3, classes=num_classes) else: model = smp.UnetPlusPlus(encoder_name=encoder_name, encoder_weights=None, in_channels=3, classes=num_classes) - - model.load_state_dict(torch.load(os.path.join(model_checkpoint, "pytorch_model.bin"), map_location=device)) + + # The Trainer saves weights as model.safetensors by default; older runs + # may have pytorch_model.bin. Load whichever is present. + safetensors_path = os.path.join(model_checkpoint, "model.safetensors") + bin_path = os.path.join(model_checkpoint, "pytorch_model.bin") + if os.path.exists(safetensors_path): + from safetensors.torch import load_file + state_dict = load_file(safetensors_path, device="cpu") + elif os.path.exists(bin_path): + state_dict = torch.load(bin_path, map_location=device) + else: + raise FileNotFoundError( + f"No SMP weights found in {model_checkpoint} " + "(looked for model.safetensors and pytorch_model.bin)." + ) + model.load_state_dict(state_dict) # SMP needs albumentations validation transform # SMP inference requires a predefined size, assuming 512 if not specified. @@ -75,29 +89,34 @@ def run_inference(model_checkpoint: str, image_path: str, output_path: str): outputs = model(**inputs) logits = outputs.logits - # 3. Post-process the logits - - # 4. Create a color mask - # We need the color map. Let's try to get it from the model config. - color_map = None - if hasattr(model.config, "id2label"): - id2label = model.config.id2label - # Create a simple color map based on label ID - # We'll just generate random colors for this demo - palette = np.random.randint(0, 255, (len(id2label), 3), dtype=np.uint8) - palette[0] = [0, 0, 0] # Make background (label 0) black - - color_map = palette[pred_seg] - color_mask = Image.fromarray(color_map.astype(np.uint8), "RGB") + # 3. Post-process the logits into a predicted segmentation map. + # HF models (e.g. SegFormer) emit logits at a reduced resolution, so we + # upsample to the original image size before taking the argmax. This also + # guarantees the mask matches the source image dimensions for overlaying. + orig_w, orig_h = image.size # PIL reports (width, height) + logits = torch.as_tensor(logits) + upsampled_logits = torch.nn.functional.interpolate( + logits, size=(orig_h, orig_w), mode="bilinear", align_corners=False + ) + pred_seg = upsampled_logits.argmax(dim=1)[0].to("cpu").numpy().astype(np.uint8) - # 5. Overlay the mask on the original image - if color_map is not None: - overlay_image = Image.blend(image, color_mask, alpha=0.5) + # 4. Build a color mask. A fixed RNG seed keeps colors stable across runs. + if is_smp: + n_labels = num_classes + elif hasattr(model, "config") and getattr(model.config, "id2label", None): + n_labels = len(model.config.id2label) else: - # Fallback: just save the raw segmentation map (will look black/white) - overlay_image = Image.fromarray(pred_seg.astype(np.uint8)) + n_labels = int(pred_seg.max()) + 1 + n_labels = max(int(n_labels), int(pred_seg.max()) + 1, 2) + + rng = np.random.default_rng(42) + palette = rng.integers(0, 256, size=(n_labels, 3), dtype=np.uint8) + palette[0] = [0, 0, 0] # background (label 0) rendered black + + color_mask = Image.fromarray(palette[pred_seg], "RGB") - # 6. Save the final image + # 5. Overlay the mask on the original image and save. + overlay_image = Image.blend(image, color_mask, alpha=0.5) overlay_image.save(output_path) print(f"Output image saved to: {output_path}") diff --git a/image_segmentation_train.py b/image_segmentation_train.py index f7cdcdb..c5902fe 100644 --- a/image_segmentation_train.py +++ b/image_segmentation_train.py @@ -10,7 +10,7 @@ import torch import glob -from datasets import load_dataset, DatasetDict, Image +from datasets import load_dataset, Dataset, DatasetDict, Image from transformers import ( Trainer, TrainingArguments, @@ -57,9 +57,36 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N labels = labels.squeeze(1) loss = loss_fct(logits, labels) - + return (loss, logits) if return_outputs else loss + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): + """ + Custom evaluation step for SMP models. + + The base Trainer calls model(**inputs), passing pixel_values= and + labels= as keywords, but an SMP model's forward() takes a single + positional image tensor. We unpack the batch and call it correctly, + returning (loss, logits, labels) in the shape the eval loop expects. + """ + inputs = self._prepare_inputs(inputs) + images = inputs["pixel_values"] + labels = inputs.get("labels") + + with torch.no_grad(): + logits = model(images) + loss = None + if labels is not None: + loss_fct = nn.CrossEntropyLoss(ignore_index=255) + target = labels + if target.dim() == 4 and target.shape[1] == 1: + target = target.squeeze(1) + loss = loss_fct(logits, target) + + if prediction_loss_only: + return (loss, None, None) + return (loss, logits, labels) + def get_trainer_class(model): """ Returns the appropriate Trainer class. @@ -146,14 +173,15 @@ def pair_images_and_masks(split_dir): # --- Step B: Load Dataset (Manual Scan or ImageFolder) --- dataset = DatasetDict() - + + # Collected across whichever loading path runs; used later for label auto-detection. + all_mask_paths = [] + if args.is_presplit: # Check standard structure splits = ["train", "validation", "test"] found_splits = [] - - all_mask_paths = [] - + for split in splits: split_dir = os.path.join(args.data_dir, split) if os.path.isdir(split_dir): @@ -331,13 +359,20 @@ def pair_images_and_masks(split_dir): ) trainer.train() - - # Explicitly save model and processor + + # Explicitly save model and processor. trainer.save_model(output_dir) - image_processor.save_pretrained(output_dir) + if image_processor is not None: + # HF models have a processor; SMP (U-Net) models do not. + image_processor.save_pretrained(output_dir) + else: + # SMP path: the Trainer saves only the state dict (no config.json), + # so we write the config ourselves. Inference relies on it to detect + # the SMP architecture and rebuild the encoder/decoder. + model.config.to_json_file(os.path.join(output_dir, "config.json")) + print(f"Wrote SMP config.json to {output_dir}") - if test_dataset is not None: print("\n--- Evaluating on Test Set ---") test_metrics = trainer.evaluate(test_dataset, metric_key_prefix="test") diff --git a/image_segmentation_utils/augmentations.py b/image_segmentation_utils/augmentations.py index 0900945..42b648f 100644 --- a/image_segmentation_utils/augmentations.py +++ b/image_segmentation_utils/augmentations.py @@ -56,7 +56,7 @@ def get_train_transform( ) transforms.append( - A.CoarseDropout(max_holes=8, max_height=32, max_width=32, min_holes=1, min_height=8, min_width=8, p=0.2) + A.CoarseDropout(num_holes_range=(1, 8), hole_height_range=(8, 32), hole_width_range=(8, 32), p=0.2) ) # 3. Normalization & Conversion (Always Applied) diff --git a/image_segmentation_utils/data_utils.py b/image_segmentation_utils/data_utils.py index 40d0a59..cff2c9b 100644 --- a/image_segmentation_utils/data_utils.py +++ b/image_segmentation_utils/data_utils.py @@ -8,21 +8,49 @@ from transformers import AutoImageProcessor from typing import List, Dict, Any +def mask_to_array(msk) -> np.ndarray: + """ + Converts a mask entry to a 2D array of class ids. + + Accepts either a decoded PIL image or the {"path", "bytes"} dict produced + by datasets' Image(decode=False). A mask's pixel values *are* the class + ids, so a palette ("P") mask must be read as raw palette indices: calling + .convert("L") on it would map the indices through the palette and silently + destroy the labels. + """ + from PIL import Image as PILImage + import io + + # Handle the undecoded {"path": ..., "bytes": ...} form. + if isinstance(msk, dict): + if msk.get("bytes") is not None: + msk = PILImage.open(io.BytesIO(msk["bytes"])) + elif msk.get("path"): + msk = PILImage.open(msk["path"]) + else: + raise ValueError("Mask dict has neither 'bytes' nor 'path'.") + + if msk.mode == "P": + return np.array(msk) + if msk.mode != "L": + msk = msk.convert("L") + return np.array(msk) + + def apply_transforms(examples: Dict[str, Any], transform: Any) -> Dict[str, Any]: """ Applies albumentations augmentations to a batch of images and masks. This function is designed to be used with dataset.map() - + Returns a dict with 'image' and 'mask' keys. """ images = [] masks = [] - + for img, msk in zip(examples["image"], examples["label"]): img = img.convert("RGB") - msk = msk.convert("L") # Ensure mask is single-channel - - transformed = transform(image=np.array(img), mask=np.array(msk)) + + transformed = transform(image=np.array(img), mask=mask_to_array(msk)) images.append(transformed["image"]) masks.append(transformed["mask"]) diff --git a/image_segmentation_utils/metrics.py b/image_segmentation_utils/metrics.py index 6874155..f953af1 100644 --- a/image_segmentation_utils/metrics.py +++ b/image_segmentation_utils/metrics.py @@ -2,6 +2,8 @@ import evaluate import numpy as np +import torch +import torch.nn.functional as F from transformers import EvalPrediction from typing import Callable, Dict, Any @@ -19,28 +21,37 @@ def compute_metrics(eval_pred: EvalPrediction) -> Dict[str, float]: for a semantic segmentation model. """ logits, labels = eval_pred.predictions, eval_pred.label_ids - + # Models like SegFormer output a tuple. # The first element [0] is the segmentation map logits. if isinstance(logits, tuple): logits = logits[0] - - # (batch_size, num_classes, height, width) -> (batch_size, height, width) - predictions = np.argmax(logits, axis=1) - # Flatten to (batch_size * height * width) - pred_flat = predictions.flatten() - label_flat = labels.flatten() + labels = np.asarray(labels) + + # SegFormer (and other HF models) emit logits at a reduced resolution + # (e.g. H/4). Upsample to the label resolution before argmax, otherwise + # predictions and references have different shapes and IoU is meaningless. + logits_t = torch.as_tensor(np.asarray(logits)) + target_hw = tuple(labels.shape[-2:]) + if tuple(logits_t.shape[-2:]) != target_hw: + logits_t = F.interpolate( + logits_t, size=target_hw, mode="bilinear", align_corners=False + ) + predictions = logits_t.argmax(dim=1).cpu().numpy() try: + # mean_iou expects lists of 2D (H, W) arrays, one per image — NOT a + # single flattened 1D array (which makes it iterate over scalars and + # raises "'numpy.int64' object has no attribute 'get'"). metrics = metric.compute( - predictions=pred_flat, - references=label_flat, + predictions=list(predictions), + references=list(labels), num_labels=num_labels, ignore_index=ignore_index, # reduce_labels=False is critical for handling batches # that don't contain every single class. - reduce_labels=False, + reduce_labels=False, ) # Format metrics diff --git a/image_segmentation_utils/model_utils.py b/image_segmentation_utils/model_utils.py index 4120cd7..cff88a9 100644 --- a/image_segmentation_utils/model_utils.py +++ b/image_segmentation_utils/model_utils.py @@ -1,7 +1,51 @@ # image_segmentation_utils/model_utils.py import segmentation_models_pytorch as smp -from transformers import AutoModelForSemanticSegmentation, AutoImageProcessor +from transformers import AutoModelForSemanticSegmentation, AutoImageProcessor, PretrainedConfig + +# Maps our SMP architecture slug to the class name we persist in config.json and +# read back at inference time. +_SMP_ARCH_TO_CLASSNAME = {"unet": "Unet", "unetplusplus": "UnetPlusPlus"} + + +def _attach_hf_config(model, arch, encoder, id2label, label2id, num_labels): + """ + Give an SMP model a Hugging Face PretrainedConfig. + + The HF Trainer assigns `model.config.use_cache = ...` in its constructor, but + SMP (>=0.5) exposes `.config` as a read-only property returning a plain dict. + We shadow that property with a per-instance subclass so `.config` becomes a + normal, settable attribute, then point it at a PretrainedConfig that both + satisfies the Trainer and carries the metadata needed to rebuild the model + at inference time. + """ + config = PretrainedConfig() + config.architectures = [_SMP_ARCH_TO_CLASSNAME.get(arch, arch)] + config.encoder_name = encoder + config.num_labels = num_labels + config.id2label = {int(k): v for k, v in id2label.items()} + config.label2id = {k: int(v) for k, v in label2id.items()} + config.use_cache = False + + # SMP's load_state_dict(self, state_dict, **kwargs) accepts `strict` only as a + # keyword, but the HF Trainer (load_best_model_at_end) passes it positionally. + # Adapt the positional call to a keyword one, preserving SMP's key remapping. + base_cls = model.__class__ + + def _load_state_dict(self, state_dict, strict=True, *args, **kwargs): + return base_cls.load_state_dict(self, state_dict, strict=strict, **kwargs) + + # Shadow the read-only `config` property with a plain class attribute (so the + # Trainer can assign model.config.use_cache) and install the load_state_dict + # adapter, both on a throwaway per-instance subclass. SMP's class is untouched. + model.__class__ = type( + f"{base_cls.__name__}HF", + (base_cls,), + {"config": None, "load_state_dict": _load_state_dict}, + ) + model.config = config + return model + def load_model(model_checkpoint, id2label, label2id): """ @@ -9,34 +53,34 @@ def load_model(model_checkpoint, id2label, label2id): Supports Hugging Face Transformers and Segmentation Models PyTorch (SMP). """ num_labels = len(id2label) - + # --- Check for SMP Models (U-Net) --- if model_checkpoint.startswith("unet"): # Format: unet-resnet34 or unetplusplus-efficientnet-b0 parts = model_checkpoint.split("-") arch = parts[0] # unet or unetplusplus encoder = "-".join(parts[1:]) # resnet34, etc. - + print(f"Loading SMP Model: Architecture={arch}, Encoder={encoder}, Classes={num_labels}") - + if arch == "unet": model = smp.Unet( - encoder_name=encoder, - encoder_weights="imagenet", - in_channels=3, + encoder_name=encoder, + encoder_weights="imagenet", + in_channels=3, classes=num_labels ) elif arch == "unetplusplus": model = smp.UnetPlusPlus( - encoder_name=encoder, - encoder_weights="imagenet", - in_channels=3, + encoder_name=encoder, + encoder_weights="imagenet", + in_channels=3, classes=num_labels ) else: raise ValueError(f"Unknown SMP architecture: {arch}") - - return model + + return _attach_hf_config(model, arch, encoder, id2label, label2id, num_labels) # --- Default: Hugging Face Transformers --- print(f"Loading Hugging Face Model: {model_checkpoint}") diff --git a/job_store.py b/job_store.py index 8a4ca0c..571a7da 100644 --- a/job_store.py +++ b/job_store.py @@ -54,21 +54,21 @@ class RegisteredDataset(Base): def init_db(): """Initialize the database and create tables if they don't exist.""" Base.metadata.create_all(bind=engine) - + # --- Database Migrations --- - # Since we added 'pid' later, we need to ensure it exists in old DBs. - # We use raw SQL for this simple migration instead of setting up Alembic. - with engine.connect() as conn: + # Since we added 'pid' later, we ensure it exists on databases created before + # that change. We inspect the schema rather than SELECTing (a raw-string + # SELECT raises ObjectNotExecutableError on SQLAlchemy 2.x) and avoid a full + # Alembic setup for a single column. + from sqlalchemy import inspect as sa_inspect, text + columns = [col["name"] for col in sa_inspect(engine).get_columns("jobs")] + if "pid" not in columns: + print("Migrating DB: Adding 'pid' column to 'jobs' table.") try: - # Try to select the column to see if it exists - conn.execute("SELECT pid FROM jobs LIMIT 1") - except Exception: - # If it fails, add the column - print("Migrating DB: Adding 'pid' column to 'jobs' table.") - try: - conn.execute("ALTER TABLE jobs ADD COLUMN pid INTEGER") - except Exception as e: - print(f"Migration warning: {e}") + with engine.begin() as conn: + conn.execute(text("ALTER TABLE jobs ADD COLUMN pid INTEGER")) + except Exception as e: + print(f"Migration warning: {e}") # --- Job Functions --- @@ -351,13 +351,15 @@ def is_valid_checkpoint_path(checkpoint_path: str) -> bool: """ if not checkpoint_path: return False - + root_path = os.path.abspath(MODEL_OUTPUT_ROOT) full_path = os.path.abspath(checkpoint_path) - - if not full_path.startswith(root_path): - return False - - return False - - return True \ No newline at end of file + + # Ensure full_path is MODEL_OUTPUT_ROOT itself or strictly inside it. + # commonpath avoids the classic startswith() prefix bug, where a sibling + # like "model_outputs_evil" would pass a "model_outputs" prefix check. + try: + return os.path.commonpath([full_path, root_path]) == root_path + except ValueError: + # Raised for paths on different drives or mixed absolute/relative. + return False \ No newline at end of file diff --git a/object_detection_train.py b/object_detection_train.py index e2ed139..4cad497 100644 --- a/object_detection_train.py +++ b/object_detection_train.py @@ -1,6 +1,5 @@ # object_detection_train.py -import os import os # Disable WandB integration os.environ["WANDB_DISABLED"] = "true" @@ -131,7 +130,7 @@ def main(args): compute_metrics=eval_compute_metrics_fn, train_dataset=dataset["train"], eval_dataset=dataset["validation"], - tokenizer=image_processor, + processing_class=image_processor, callbacks=[json_metrics_callback, early_stopping_callback] ) diff --git a/requirements.txt b/requirements.txt index 7f9995a..7b7ee79 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,40 +1,54 @@ +# Pinned to a verified-working stack. +# +# Every version below was confirmed by running the full train -> evaluate -> +# save -> inference cycle for all three tasks and both model families per task +# (image classification, object detection via DETR and YOLO, semantic +# segmentation via SegFormer and SMP U-Net) on Python 3.12 / CPU torch. +# +# Torch is pinned by version only (no +cpu local tag): on the CUDA base image in +# the Dockerfile, PyPI serves the matching CUDA build of the same version. +# +# To refresh: bump a pin, re-run the smoke tests for all tasks, then commit. + # --- Core ML/DL --- -torch -torchvision -torchaudio -accelerate +torch==2.13.0 +torchvision==0.28.0 +torchaudio==2.11.0 +accelerate==1.14.0 # --- Hugging Face --- -transformers -datasets -evaluate -huggingface_hub +transformers==5.13.1 +datasets==5.0.0 +evaluate==0.4.6 +huggingface_hub==1.23.0 # --- Task-Specific --- -albumentations -torchmetrics -timm -pycocotools -ultralytics -faster-coco-eval -segmentation-models-pytorch +albumentations==2.0.8 +torchmetrics==1.9.0 +timm==1.0.28 +pycocotools==2.0.11 +ultralytics==8.4.92 +faster-coco-eval==1.7.2 +segmentation-models-pytorch==0.5.0 -# --- API & UI --- -fastapi -uvicorn[standard] -python-multipart -gradio +# --- API --- +fastapi==0.139.0 +uvicorn[standard]==0.51.0 +python-multipart==0.0.32 # --- Utilities --- -numpy -pandas -scipy -Pillow -wandb -psutil -scikit-learn +numpy==2.4.4 +pandas==3.0.3 +scipy==1.18.0 +Pillow==12.2.0 +psutil==7.2.2 +scikit-learn==1.9.0 # --- Job Persistence & Queuing --- -sqlalchemy -celery -redis \ No newline at end of file +sqlalchemy==2.0.51 +celery==5.6.3 +redis==8.0.1 + +# Removed (2026-07): gradio and wandb — neither is imported anywhere in the +# codebase (training sets WANDB_DISABLED=true), and gradio pulls a large, +# unused dependency tree. Re-add with a pin if a demo UI / W&B logging is built. diff --git a/scripts/smoke_test.py b/scripts/smoke_test.py new file mode 100644 index 0000000..163869d --- /dev/null +++ b/scripts/smoke_test.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +"""End-to-end smoke test for the deep-learning pipelines. + +Generates tiny synthetic datasets and runs the real training and inference +scripts for every task and model family, asserting each completes and produces +the expected artifacts. Designed to run on CPU in CI as a regression guard: it +is what catches library-version drift before it reaches production. + +Run: python scripts/smoke_test.py +Env overrides (used by CI to pick small models): + SMOKE_CLS_MODEL, SMOKE_DET_HF_MODEL, SMOKE_DET_YOLO_MODEL, + SMOKE_SEG_HF_MODEL, SMOKE_SEG_SMP_MODEL + SMOKE_KEEP=1 keep the work dir for inspection +""" +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np +from PIL import Image + +REPO = Path(__file__).resolve().parent.parent +PY = sys.executable + +CLS_MODEL = os.environ.get("SMOKE_CLS_MODEL", "google/efficientnet-b0") +DET_HF_MODEL = os.environ.get("SMOKE_DET_HF_MODEL", "facebook/detr-resnet-50") +DET_YOLO_MODEL = os.environ.get("SMOKE_DET_YOLO_MODEL", "yolo11n.pt") +SEG_HF_MODEL = os.environ.get("SMOKE_SEG_HF_MODEL", "nvidia/segformer-b0-finetuned-ade-512-512") +SEG_SMP_MODEL = os.environ.get("SMOKE_SEG_SMP_MODEL", "unet-resnet18") + +rng = np.random.default_rng(0) + + +# --------------------------------------------------------------------------- # +# Fixture generation (synthetic data only) +# --------------------------------------------------------------------------- # +def make_classification(root: Path): + for split, n in {"train": 6, "validation": 2, "test": 2}.items(): + for cls, (lo, hi) in {"dark": (0, 80), "bright": (150, 256)}.items(): + d = root / split / cls + d.mkdir(parents=True, exist_ok=True) + for i in range(n): + arr = rng.integers(lo, hi, (64, 64, 3), dtype=np.uint8) + Image.fromarray(arr, "RGB").save(d / f"{cls}_{i}.png") + + +def make_detection(root: Path): + import json + for split, n in {"train": 6, "validation": 2, "test": 2}.items(): + d = root / split + d.mkdir(parents=True, exist_ok=True) + images, annotations, ann_id = [], [], 1 + for img_id in range(n): + img = rng.integers(0, 60, (96, 96, 3), dtype=np.uint8) + bw, bh = int(rng.integers(18, 30)), int(rng.integers(18, 30)) + x, y = int(rng.integers(0, 96 - bw)), int(rng.integers(0, 96 - bh)) + img[y:y + bh, x:x + bw] = rng.integers(180, 256, 3, dtype=np.uint8) + fname = f"img_{img_id}.png" + Image.fromarray(img, "RGB").save(d / fname) + images.append({"id": img_id, "file_name": fname, "width": 96, "height": 96}) + annotations.append({"id": ann_id, "image_id": img_id, "category_id": 1, + "bbox": [x, y, bw, bh], "area": bw * bh, "iscrowd": 0}) + ann_id += 1 + (d / "_annotations.coco.json").write_text(json.dumps( + {"images": images, "annotations": annotations, + "categories": [{"id": 1, "name": "object", "supercategory": "none"}]})) + + +def make_segmentation(root: Path): + for split, n in {"train": 6, "validation": 2, "test": 2}.items(): + (root / split / "images").mkdir(parents=True, exist_ok=True) + (root / split / "masks").mkdir(parents=True, exist_ok=True) + for i in range(n): + img = rng.integers(0, 90, (64, 64, 3), dtype=np.uint8) + mask = np.zeros((64, 64), dtype=np.uint8) + y, x = int(rng.integers(8, 36)), int(rng.integers(8, 36)) + img[y:y + 18, x:x + 18] = rng.integers(170, 256, 3, dtype=np.uint8) + mask[y:y + 18, x:x + 18] = 1 + Image.fromarray(img, "RGB").save(root / split / "images" / f"s_{i}.png") + Image.fromarray(mask, "L").save(root / split / "masks" / f"s_{i}.png") + + +# --------------------------------------------------------------------------- # +# Runner helpers +# --------------------------------------------------------------------------- # +def run(cmd, env): + print(" $ " + " ".join(str(c) for c in cmd), flush=True) + p = subprocess.run(cmd, cwd=REPO, env=env, capture_output=True, text=True) + if p.returncode != 0: + print(p.stdout[-3000:]); print(p.stderr[-3000:]) + return p.returncode == 0 + + +def find_run_dir(root: Path, prefix: str): + hits = [d for d in root.glob(f"{prefix}*") if d.is_dir()] + return hits[0] if hits else None + + +# --------------------------------------------------------------------------- # +# Task smoke tests — each returns (name, ok) +# --------------------------------------------------------------------------- # +def smoke_classification(work: Path, env): + data, out = work / "cls_data", work / "cls_out" + make_classification(data) + ok = run([PY, "image_classification_train.py", "--data_dir", data, + "--model_checkpoint", CLS_MODEL, "--is_presplit", "--epochs", "1", + "--train_batch_size", "4", "--eval_batch_size", "4", "--max_image_size", "224", + "--num_proc", "1", "--model_output_root", out, "--metrics_filename", "m.json", + "--run_name", "smoke", "--version", "0"], env) + ckpt = find_run_dir(out, CLS_MODEL.split("/")[-1]) + if ok and ckpt: + img = next((data / "test").rglob("*.png")) + ok = run([PY, "image_classification_inference.py", + "--model_checkpoint", ckpt, "--image_path", img], env) + return ("classification", ok and ckpt is not None) + + +def smoke_detection_hf(work: Path, env): + data, proc, out = work / "det_data", work / "det_proc", work / "det_hf_out" + make_detection(data) + ok = run([PY, "object_detection_train.py", "--data_dir", data, "--processed_data_dir", proc, + "--model_checkpoint", DET_HF_MODEL, "--epochs", "1", "--train_batch_size", "2", + "--eval_batch_size", "2", "--max_image_size", "256", "--num_proc", "1", + "--model_output_root", out, "--metrics_filename", "m.json", + "--run_name", "smoke", "--version", "0"], env) + ckpt = find_run_dir(out, DET_HF_MODEL.split("/")[-1]) + if ok and ckpt: + img = next((data / "test").glob("*.png")) + ok = run([PY, "object_detection_inference.py", "--model_checkpoint", ckpt, + "--image_path", img, "--output_path", work / "det_hf_pred.png", + "--threshold", "0.1"], env) + return ("detection-DETR", ok and ckpt is not None) + + +def smoke_detection_yolo(work: Path, env): + data, proc, out = work / "det_data", work / "det_proc", work / "det_yolo_out" + if not proc.exists(): # reuse the Arrow dataset the DETR run built + make_detection(data) + subprocess.run([PY, "object_detection_utils/preprocess_data.py", + "--raw_data_dir", data, "--processed_data_dir", proc], + cwd=REPO, env=env, capture_output=True, text=True) + run_dir = out / "yolo-run" + ok = run([PY, "object_detection_train_yolo.py", "--data_dir", proc, + "--model_checkpoint", DET_YOLO_MODEL, "--run_name", "smoke", "--version", "0", + "--model_output_root", out, "--metrics_filename", "m.json", + "--output_dir", run_dir, "--epochs", "2", "--train_batch_size", "2", + "--max_image_size", "96", "--num_proc", "1", "--early_stopping_patience", "5"], env) + best = run_dir / "weights" / "best.pt" + if ok and best.exists(): + img = next((data / "test").glob("*.png")) + ok = run([PY, "object_detection_inference_yolo.py", "--model_checkpoint", run_dir, + "--image_path", img, "--output_path", work / "det_yolo_pred.png", + "--threshold", "0.1", "--iou", "0.7", "--max_det", "300", "--imgsz", "96"], env) + return ("detection-YOLO", ok and best.exists()) + + +def smoke_segmentation(work: Path, env, model, tag): + data, out = work / "seg_data", work / f"seg_{tag}_out" + make_segmentation(data) + ok = run([PY, "image_segmentation_train.py", "--data_dir", data, + "--model_checkpoint", model, "--is_presplit", "--epochs", "1", + "--train_batch_size", "2", "--eval_batch_size", "2", "--max_image_size", "64", + "--num_proc", "1", "--model_output_root", out, "--metrics_filename", "m.json", + "--run_name", tag, "--version", "0"], env) + ckpt = find_run_dir(out, model.split("/")[-1]) + if ok and ckpt: + img = next((data / "test" / "images").glob("*.png")) + ok = run([PY, "image_segmentation_inference.py", "--model_checkpoint", ckpt, + "--image_path", img, "--output_path", work / f"seg_{tag}_pred.png"], env) + return (f"segmentation-{tag}", ok and ckpt is not None) + + +def main(): + env = os.environ.copy() + env.update({"WANDB_DISABLED": "true", "TQDM_DISABLE": "1", + "TOKENIZERS_PARALLELISM": "false", + "HF_HOME": os.environ.get("HF_HOME", str(Path.home() / ".cache/huggingface"))}) + work = Path(tempfile.mkdtemp(prefix="dl_smoke_")) + print(f"Work dir: {work}\n") + + results = [] + for label, fn in [ + ("classification", lambda: smoke_classification(work, env)), + ("detection-DETR", lambda: smoke_detection_hf(work, env)), + ("detection-YOLO", lambda: smoke_detection_yolo(work, env)), + ("segmentation-SegFormer", lambda: smoke_segmentation(work, env, SEG_HF_MODEL, "sf")), + ("segmentation-UNet", lambda: smoke_segmentation(work, env, SEG_SMP_MODEL, "unet")), + ]: + print(f"=== {label} ===", flush=True) + try: + results.append(fn()) + except Exception as e: + print(f" EXCEPTION: {e}") + results.append((label, False)) + print() + + if not os.environ.get("SMOKE_KEEP"): + import shutil + shutil.rmtree(work, ignore_errors=True) + + print("=" * 40) + all_ok = True + for name, ok in results: + print(f" {'PASS' if ok else 'FAIL'} {name}") + all_ok = all_ok and ok + print("=" * 40) + sys.exit(0 if all_ok else 1) + + +if __name__ == "__main__": + main()