diff --git a/README.md b/README.md
index a405ce0..8debbff 100644
--- a/README.md
+++ b/README.md
@@ -59,6 +59,8 @@ project/
├── 2_train_classifier.py ← train EfficientNet classifier
├── 3_inference.py ← end-to-end detect + classify on new images
├── api_server.py ← FastAPI polling server
+├── balance_and_augment.py ← balance classes + generate augmented images
+├── kaggle_train.py ← automate Kaggle upload + training + download
└── train_classifier_kaggle.ipynb
```
@@ -311,6 +313,75 @@ The splitter prints a table so you can spot imbalanced or thin classes before tr
---
+### Step 3b — Balance & augment the dataset
+
+> Optional but strongly recommended when classes have unequal image counts or any class has fewer than 50 training images.
+
+Analyses the `train/` split, calculates how many images each class needs to reach a target count, then generates augmented copies using a three-tier albumentations pipeline until every class is balanced.
+
+```bash
+python balance_and_augment.py --data_dir crops_dataset --split train --max_scale 2.0 --workers 8 --dry_run
+```
+
+Remove `--dry_run` once you are happy with the plan.
+
+| Argument | Default | Description |
+|---|---|---|
+| `--data_dir` | `crops_dataset` | Root with `train/` `val/` `test/` subfolders |
+| `--split` | `train` | Which split to augment |
+| `--target_count` | auto | Explicit target per class — omit to use auto |
+| `--max_scale` | `2.0` | Auto target = largest class × this factor |
+| `--min_count` | `0.95` | Skip classes already within 5 % of target |
+| `--workers` | cpu − 1 | Parallel workers |
+| `--dry_run` | off | Print the plan without writing files |
+| `--val_split` | `0.0` | Fraction of augmented images also copied to `val/` |
+| `--suffix` | `_aug` | Filename suffix added to augmented images |
+| `--quality` | `92` | JPEG save quality |
+
+**Dry-run output example:**
+
+```
+──────────────────────────────────────────────────────────
+ Class Current Target To add Status
+──────────────────────────────────────────────────────────
+ cola_can 72 200 128 + 128 ████████████
+ lays_chips 18 200 182 + 182 ██
+ pepsi_can 58 200 142 + 142 ██████
+ siglo 204 200 0 ✓ ok ████████████████████
+──────────────────────────────────────────────────────────
+ TOTAL 352 452
+──────────────────────────────────────────────────────────
+```
+
+**Three-tier augmentation pipeline:**
+
+| Tier | Transforms | Probability |
+|---|---|---|
+| **Light** (always) | HorizontalFlip, ShiftScaleRotate ±15°, BrightnessContrast ±30%, HueSaturation | 75–85% each |
+| **Medium** (random subset) | Perspective, GridDistortion, ElasticTransform, Blur/Sharpen, GaussNoise, CoarseDropout (occlusion), JPEG compression | 25–40% |
+| **Heavy** (rare) | RandomShadow, RandomSunFlare, RandomFog, ChannelShuffle, RGBShift | 8–15% |
+
+Augmented files are written **beside the originals** in the same class folder. Originals are never modified or deleted.
+
+> **When to use `--val_split`:** Pass `--val_split 0.10` to also copy 10% of the augmented images into `val/` if your val set is very small (< 5 images per class). For most cases, leave it at `0` — val should represent real unaugmented images.
+
+---
+
+### Step 3b — recommended workflow
+
+```bash
+# 1. Preview the plan
+python balance_and_augment.py --data_dir crops_dataset --dry_run
+
+# 2. Run with auto target (largest class × 2)
+python balance_and_augment.py --data_dir crops_dataset --workers 8
+
+# 3. Or set an explicit target
+python balance_and_augment.py --data_dir crops_dataset --target_count 300
+```
+
+---
+
### Step 4 — Train the classifier
Trains an **EfficientNet-B0** (ImageNet pre-trained) on `crops_dataset/`. A `WeightedRandomSampler` is used automatically so class imbalance does not bias training.
@@ -399,6 +470,72 @@ After training, download from `/kaggle/working/runs/classify/`:
These two files are all you need to run `3_inference.py` and `api_server.py`.
+#### Option C — Fully automated via `kaggle_train.py`
+
+Run the entire Kaggle pipeline — upload, train, download — from a single command with no browser needed:
+
+```bash
+python kaggle_train.py \
+ --dataset_dir crops_dataset \
+ --notebook train_classifier_kaggle.ipynb \
+ --dataset_name my-product-crops \
+ --kernel_name product-classifier-train \
+ --output_dir runs/classify
+```
+
+**One-time setup:**
+
+1. Go to [kaggle.com/settings/api](https://www.kaggle.com/settings/api) → **Create New Token** → downloads `kaggle.json`
+2. Move it to `~/.kaggle/kaggle.json`
+3. On Linux/Mac: `chmod 600 ~/.kaggle/kaggle.json`
+
+Or use environment variables instead of the file:
+```bash
+export KAGGLE_USERNAME=your_username
+export KAGGLE_KEY=your_api_key
+```
+
+**What the script does automatically:**
+
+| Step | Action |
+|---|---|
+| 1 | Zips `crops_dataset/` and uploads (or updates) it as a Kaggle Dataset |
+| 2 | Patches `DATA_DIR` in the notebook to match the uploaded dataset path |
+| 3 | Pushes the notebook as a Kaggle Kernel and triggers a GPU run |
+| 4 | Polls kernel status every 30 s until `complete` or `failed` |
+| 5 | Downloads `best.pt`, `class_names.json`, and plot PNGs into `runs/classify/` |
+
+**Arguments:**
+
+| Argument | Default | Description |
+|---|---|---|
+| `--dataset_dir` | `crops_dataset` | Local dataset folder to upload |
+| `--notebook` | `train_classifier_kaggle.ipynb` | Notebook file to push |
+| `--dataset_name` | `product-crops` | Kaggle dataset slug (lowercase, hyphens) |
+| `--kernel_name` | `product-classifier` | Kaggle kernel slug |
+| `--output_dir` | `runs/classify` | Where to save downloaded weights |
+| `--poll_interval` | `30` | Seconds between status checks |
+| `--timeout` | `180` | Max minutes to wait before giving up |
+| `--skip_upload` | off | Skip zip+upload — reuse existing Kaggle dataset |
+| `--skip_push` | off | Skip kernel push — just poll and download last run |
+| `--no_gpu` | off | Run on CPU instead of T4 GPU |
+| `--public` | off | Make dataset and kernel public |
+
+**Useful combinations:**
+
+```bash
+# Re-run training on an already uploaded dataset (faster — skips the zip/upload)
+python kaggle_train.py --skip_upload --dataset_name my-product-crops --kernel_name product-classifier-train
+
+# Just download the outputs of the last completed run (no upload, no push)
+python kaggle_train.py --skip_upload --skip_push --kernel_name product-classifier-train
+
+# Full run but with a fresh public kernel
+python kaggle_train.py --dataset_name my-product-crops --kernel_name product-classifier-train --public
+```
+
+---
+
#### Diagnosing a stuck val loss
A val loss stuck above `ln(num_classes)` (e.g. > 3.13 for 23 classes) means the model is guessing randomly or confidently wrong. Common causes:
@@ -452,10 +589,8 @@ The classifier automatically probes available VRAM at startup and picks the larg
---
## 5. API Server
+
A REST API with a polling pattern for integration into other applications.
-
-
-
```bash
pip install fastapi uvicorn
@@ -560,6 +695,8 @@ Progress stages:
| `1_generate_crop_dataset.py` | Detect & crop (no classifier) | shelf images | `data/unknown/*.jpg` |
| `auto_label.py` | Detect + auto-classify + browser review | shelf images + both models | `data//*.jpg` |
| `1b_split_dataset.py` | Stratified train/val/test split | `data/` | `crops_dataset/` |
+| `balance_and_augment.py` | Balance classes + augment | `crops_dataset/train/` | augmented images in-place |
+| `kaggle_train.py` | Automate Kaggle upload + train + download | `crops_dataset/` + notebook | `runs/classify/best.pt` |
| `2_train_classifier.py` | Train EfficientNet classifier | `crops_dataset/` | `runs/classify/best.pt` |
| `3_inference.py` | End-to-end detect + classify | images / video / webcam | annotated images/video |
| `api_server.py` | REST polling API | — | JSON + base64 annotated image |
@@ -620,7 +757,9 @@ Review UI (confirm / fix / reject)
↓
1b_split_dataset.py (re-split the grown dataset)
↓
-2_train_classifier.py (retrain from scratch or from last.pt)
+balance_and_augment.py (equalise class counts, generate augmented copies)
+ ↓
+2_train_classifier.py (local) OR kaggle_train.py (automated Kaggle)
↓
3_inference.py / api_server.py (deploy updated model)
↓
diff --git a/kaggle_train.py b/kaggle_train.py
new file mode 100644
index 0000000..c37d8bd
--- /dev/null
+++ b/kaggle_train.py
@@ -0,0 +1,536 @@
+"""
+kaggle_train.py — Automate Kaggle Training
+============================================
+Full automation of the Kaggle training pipeline:
+
+ Step 1 Zip crops_dataset/ and upload (or update) it as a Kaggle Dataset
+ Step 2 Patch the notebook's DATA_DIR to match the uploaded dataset path
+ Step 3 Push the notebook as a Kaggle Kernel and trigger a run
+ Step 4 Poll the kernel status every 30 s until complete or failed
+ Step 5 Download best.pt + class_names.json into your local runs/classify/
+
+Requirements
+------------
+ pip install kaggle
+
+ Set up credentials (one-time):
+ Go to https://www.kaggle.com/settings/api → "Create New Token"
+ Save the downloaded kaggle.json to ~/.kaggle/kaggle.json
+ chmod 600 ~/.kaggle/kaggle.json # Linux/Mac only
+
+Usage
+-----
+# First time — creates the Kaggle dataset and notebook from scratch
+python kaggle_train.py \
+ --dataset_dir crops_dataset \
+ --notebook train_classifier_kaggle.ipynb \
+ --dataset_name my-product-crops \
+ --kernel_name product-classifier-train \
+ --output_dir runs/classify
+
+# Subsequent runs — detects existing dataset/kernel and updates them
+python kaggle_train.py \
+ --dataset_dir crops_dataset \
+ --notebook train_classifier_kaggle.ipynb \
+ --dataset_name my-product-crops \
+ --kernel_name product-classifier-train
+
+Arguments
+---------
+--dataset_dir Local crops_dataset/ folder to upload (default: crops_dataset)
+--notebook Kaggle notebook .ipynb file (default: train_classifier_kaggle.ipynb)
+--dataset_name Kaggle dataset slug (lowercase, hyphens) (default: product-crops)
+--kernel_name Kaggle kernel slug (lowercase, hyphens) (default: product-classifier)
+--output_dir Where to save downloaded weights (default: runs/classify)
+--poll_interval Seconds between status checks (default: 30)
+--timeout Max minutes to wait for kernel (default: 180)
+--skip_upload Skip dataset upload (use existing version) (flag)
+--skip_push Skip kernel push (use last run) (flag)
+--gpu Request GPU accelerator (default: True)
+--public Make dataset/kernel public (default: False)
+"""
+
+import argparse
+import json
+import os
+import shutil
+import sys
+import tempfile
+import time
+import zipfile
+from pathlib import Path
+
+
+# ─────────────────────────── args ────────────────────────────────────────────
+
+def parse_args():
+ p = argparse.ArgumentParser(
+ description="Automate Kaggle dataset upload + notebook run + output download",
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+ )
+ p.add_argument("--dataset_dir", default="crops_dataset")
+ p.add_argument("--notebook", default="train_classifier_kaggle.ipynb")
+ p.add_argument("--dataset_name", default="product-crops",
+ help="Kaggle dataset slug — lowercase, hyphens, no spaces")
+ p.add_argument("--kernel_name", default="product-classifier",
+ help="Kaggle kernel slug — lowercase, hyphens, no spaces")
+ p.add_argument("--output_dir", default="runs/classify")
+ p.add_argument("--poll_interval", type=int, default=30,
+ help="Seconds between kernel status polls")
+ p.add_argument("--timeout", type=int, default=180,
+ help="Max minutes to wait for the kernel to finish")
+ p.add_argument("--skip_upload", action="store_true",
+ help="Skip dataset upload — use the existing Kaggle dataset version")
+ p.add_argument("--skip_push", action="store_true",
+ help="Skip kernel push — just poll and download last run")
+ p.add_argument("--gpu", action="store_true", default=True,
+ help="Request GPU (T4) accelerator for the kernel")
+ p.add_argument("--no_gpu", action="store_true",
+ help="Override --gpu: run on CPU instead")
+ p.add_argument("--public", action="store_true",
+ help="Make dataset and kernel public (default: private)")
+ return p.parse_args()
+
+
+# ─────────────────────────── helpers ─────────────────────────────────────────
+
+def log(msg: str, icon: str = "▸"):
+ ts = time.strftime("%H:%M:%S")
+ print(f"[{ts}] {icon} {msg}", flush=True)
+
+def fatal(msg: str):
+ print(f"\n[ERROR] {msg}", file=sys.stderr)
+ sys.exit(1)
+
+def check_credentials():
+ """Verify kaggle.json exists before doing anything."""
+ cred_path = Path.home() / ".kaggle" / "kaggle.json"
+ env_user = os.getenv("KAGGLE_USERNAME")
+ env_key = os.getenv("KAGGLE_KEY")
+
+ if not cred_path.exists() and not (env_user and env_key):
+ fatal(
+ "Kaggle credentials not found.\n\n"
+ "Option 1 (recommended):\n"
+ " 1. Go to https://www.kaggle.com/settings/api\n"
+ " 2. Click 'Create New Token' → saves kaggle.json\n"
+ " 3. Move it to ~/.kaggle/kaggle.json\n"
+ " 4. chmod 600 ~/.kaggle/kaggle.json (Linux/Mac)\n\n"
+ "Option 2 (environment variables):\n"
+ " export KAGGLE_USERNAME=your_username\n"
+ " export KAGGLE_KEY=your_api_key\n"
+ )
+ log("Kaggle credentials found", "✓")
+
+def get_kaggle_username() -> str:
+ """Read username from kaggle.json or environment variable."""
+ env = os.getenv("KAGGLE_USERNAME")
+ if env:
+ return env
+ cred = Path.home() / ".kaggle" / "kaggle.json"
+ return json.loads(cred.read_text())["username"]
+
+def slugify(name: str) -> str:
+ """Ensure a slug is lowercase with hyphens."""
+ return name.lower().replace("_", "-").replace(" ", "-")
+
+
+# ─────────────────────────── step 1: zip + upload dataset ────────────────────
+
+def zip_dataset(dataset_dir: Path, tmp_dir: Path) -> Path:
+ """Zip the dataset folder, preserving the train/val/test structure."""
+ zip_path = tmp_dir / "dataset.zip"
+ log(f"Zipping {dataset_dir} …", "📦")
+
+ total = sum(1 for p in dataset_dir.rglob("*") if p.is_file())
+ done = 0
+
+ with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+ for file_path in sorted(dataset_dir.rglob("*")):
+ if file_path.is_file():
+ arcname = file_path.relative_to(dataset_dir.parent)
+ zf.write(file_path, arcname)
+ done += 1
+ if done % 500 == 0 or done == total:
+ pct = done / total * 100
+ print(f" Zipped {done:,}/{total:,} ({pct:.0f}%)", end="\r")
+
+ size_mb = zip_path.stat().st_size / 1e6
+ print()
+ log(f"Zip ready: {zip_path} ({size_mb:.1f} MB)", "✓")
+ return zip_path
+
+
+def upload_dataset(api, username: str, dataset_name: str,
+ dataset_dir: Path, tmp_dir: Path, public: bool) -> str:
+ """
+ Create or update a Kaggle dataset from dataset_dir.
+ Returns the full dataset reference string: username/dataset-name
+ """
+ slug = slugify(dataset_name)
+ ref = f"{username}/{slug}"
+
+ # Check if dataset already exists
+ existing = None
+ try:
+ existing = api.dataset_status(username, slug)
+ except Exception:
+ pass # 404 = doesn't exist yet
+
+ # Write dataset-metadata.json into a staging folder
+ stage_dir = tmp_dir / "dataset_stage"
+ stage_dir.mkdir(exist_ok=True)
+
+ # Copy the crops_dataset folder into staging
+ dst = stage_dir / dataset_dir.name
+ if dst.exists():
+ shutil.rmtree(dst)
+ shutil.copytree(str(dataset_dir), str(dst))
+
+ meta = {
+ "title": dataset_name.replace("-", " ").title(),
+ "id": ref,
+ "licenses": [{"name": "CC0-1.0"}],
+ }
+ (stage_dir / "dataset-metadata.json").write_text(json.dumps(meta, indent=2))
+
+ if existing is None:
+ log(f"Creating new Kaggle dataset: {ref}", "☁")
+ api.dataset_create_new(
+ folder = str(stage_dir),
+ public = public,
+ quiet = False,
+ convert_to_csv = False,
+ dir_mode = "zip",
+ )
+ log(f"Dataset created: https://www.kaggle.com/datasets/{ref}", "✓")
+ else:
+ log(f"Updating existing Kaggle dataset: {ref}", "☁")
+ api.dataset_create_version(
+ folder = str(stage_dir),
+ version_notes = f"Updated via kaggle_train.py at {time.strftime('%Y-%m-%d %H:%M')}",
+ quiet = False,
+ convert_to_csv = False,
+ delete_old_versions = False,
+ dir_mode = "zip",
+ )
+ log(f"Dataset updated: https://www.kaggle.com/datasets/{ref}", "✓")
+
+ return ref
+
+
+# ─────────────────────────── step 2: patch notebook ──────────────────────────
+
+def patch_notebook(notebook_path: Path, dataset_ref: str,
+ tmp_dir: Path) -> Path:
+ """
+ Patch DATA_DIR in the notebook to point to the uploaded Kaggle dataset,
+ and save a modified copy into tmp_dir.
+
+ The notebook's Cell 2 contains something like:
+ DATA_DIR = Path("/kaggle/input/your-dataset-name")
+ We replace the path to match the actual uploaded dataset slug.
+ """
+ nb = json.loads(notebook_path.read_text())
+
+ # Extract the dataset slug from "username/slug" → "slug"
+ ds_slug = dataset_ref.split("/")[-1]
+
+ patched = False
+ for cell in nb.get("cells", []):
+ if cell.get("cell_type") != "code":
+ continue
+ src = "".join(cell["source"])
+ if "DATA_DIR" in src and "/kaggle/input/" in src:
+ new_lines = []
+ for line in cell["source"]:
+ if "DATA_DIR" in line and "/kaggle/input/" in line:
+ # Preserve indentation, replace only the path string
+ indent = len(line) - len(line.lstrip())
+ new_line = " " * indent + f'DATA_DIR = Path("/kaggle/input/{ds_slug}/{Path(ds_slug).stem}")\n'
+ # Try to find the inner folder name (crops_dataset by default)
+ # Use the dataset_dir name if we can detect it from context
+ new_line = " " * indent + f'DATA_DIR = Path("/kaggle/input/{ds_slug}")\n'
+ new_lines.append(new_line)
+ patched = True
+ log(f"Patched DATA_DIR → /kaggle/input/{ds_slug}", "✓")
+ else:
+ new_lines.append(line)
+ cell["source"] = new_lines
+
+ if not patched:
+ log("⚠ DATA_DIR line not found in notebook — using notebook as-is", "⚠")
+
+ out_path = tmp_dir / notebook_path.name
+ out_path.write_text(json.dumps(nb, indent=1))
+ return out_path
+
+
+# ─────────────────────────── step 3: push kernel ─────────────────────────────
+
+def push_kernel(api, username: str, kernel_name: str, notebook_path: Path,
+ dataset_ref: str, tmp_dir: Path,
+ gpu: bool, public: bool) -> str:
+ """
+ Create or update a Kaggle kernel (notebook) and trigger a run.
+ Returns the kernel ref string: username/kernel-name
+ """
+ slug = slugify(kernel_name)
+ ref = f"{username}/{slug}"
+
+ kernel_dir = tmp_dir / "kernel_push"
+ kernel_dir.mkdir(exist_ok=True)
+
+ # Copy patched notebook into kernel dir
+ nb_dest = kernel_dir / notebook_path.name
+ shutil.copy2(str(notebook_path), str(nb_dest))
+
+ # kernel-metadata.json — controls GPU, dataset attachment, language
+ meta = {
+ "id": ref,
+ "title": kernel_name.replace("-", " ").title(),
+ "code_file": notebook_path.name,
+ "language": "python",
+ "kernel_type": "notebook",
+ "is_private": not public,
+ "enable_gpu": gpu,
+ "enable_tpu": False,
+ "enable_internet": True,
+ "dataset_sources": [dataset_ref],
+ "competition_sources": [],
+ "kernel_sources": [],
+ "model_sources": [],
+ }
+ (kernel_dir / "kernel-metadata.json").write_text(json.dumps(meta, indent=2))
+
+ log(f"Pushing kernel to Kaggle: {ref} (GPU={'yes' if gpu else 'no'})", "🚀")
+ api.kernels_push(str(kernel_dir))
+ log(f"Kernel pushed — run started: https://www.kaggle.com/code/{ref}", "✓")
+ return ref
+
+
+# ─────────────────────────── step 4: poll ────────────────────────────────────
+
+def poll_kernel(api, kernel_ref: str, poll_interval: int, timeout_minutes: int) -> bool:
+ """
+ Poll the kernel until it reaches a terminal status.
+ Returns True if successful, False if failed or timed out.
+ """
+ username, slug = kernel_ref.split("/")
+ deadline = time.time() + timeout_minutes * 60
+ attempt = 0
+
+ TERMINAL = {"complete", "error", "cancelled"}
+ RUNNING = {"running", "queued"}
+
+ log(f"Polling kernel status every {poll_interval}s "
+ f"(timeout {timeout_minutes} min) …", "⏳")
+
+ while time.time() < deadline:
+ attempt += 1
+ try:
+ status_obj = api.kernel_status(username, slug)
+ # The API returns an object with a .status attribute
+ status = getattr(status_obj, "status", str(status_obj)).lower()
+ except Exception as e:
+ log(f"Poll error (will retry): {e}", "⚠")
+ time.sleep(poll_interval)
+ continue
+
+ elapsed_min = (time.time() + timeout_minutes * 60 - deadline) / 60 + timeout_minutes
+ print(f" [{attempt:>3}] status={status:<12} elapsed={elapsed_min:.1f} min",
+ end="\r", flush=True)
+
+ if status == "complete":
+ print()
+ log("Kernel completed successfully!", "✓")
+ return True
+
+ if status == "error":
+ print()
+ log("Kernel run failed — check logs at "
+ f"https://www.kaggle.com/code/{kernel_ref}", "✗")
+ return False
+
+ if status == "cancelled":
+ print()
+ log("Kernel was cancelled.", "✗")
+ return False
+
+ if status not in RUNNING:
+ print()
+ log(f"Unknown status: {status!r} — continuing to poll", "⚠")
+
+ time.sleep(poll_interval)
+
+ print()
+ log(f"Timeout after {timeout_minutes} minutes. "
+ f"Check manually: https://www.kaggle.com/code/{kernel_ref}", "✗")
+ return False
+
+
+# ─────────────────────────── step 5: download outputs ────────────────────────
+
+def download_outputs(api, kernel_ref: str, output_dir: Path) -> list[Path]:
+ """
+ Download all kernel output files and return paths to the ones we care about.
+ Targets: best.pt, last.pt, class_names.json, *.png
+ """
+ username, slug = kernel_ref.split("/")
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ log(f"Downloading kernel outputs → {output_dir}", "⬇")
+
+ with tempfile.TemporaryDirectory() as tmp:
+ try:
+ api.kernels_output(username, slug, path=tmp, unzip=True)
+ except Exception as e:
+ fatal(f"Could not download kernel outputs: {e}\n"
+ f"Try manually at https://www.kaggle.com/code/{kernel_ref}/output")
+
+ tmp_path = Path(tmp)
+ downloaded = []
+
+ # Walk everything the kernel produced
+ for f in sorted(tmp_path.rglob("*")):
+ if not f.is_file():
+ continue
+
+ # We want weights, class map, and training plots
+ keep = f.suffix in {".pt", ".json", ".png", ".csv"}
+ if not keep:
+ continue
+
+ dst = output_dir / f.relative_to(tmp_path)
+ dst.parent.mkdir(parents=True, exist_ok=True)
+ shutil.copy2(str(f), str(dst))
+ downloaded.append(dst)
+ log(f" {dst.relative_to(output_dir)} ({f.stat().st_size/1e6:.2f} MB)", " ")
+
+ if not downloaded:
+ log("No output files found. The kernel may not have saved to /kaggle/working/", "⚠")
+ else:
+ log(f"{len(downloaded)} file(s) downloaded to {output_dir.resolve()}", "✓")
+
+ return downloaded
+
+
+# ─────────────────────────── main ────────────────────────────────────────────
+
+def main():
+ args = parse_args()
+
+ gpu = args.gpu and not args.no_gpu
+
+ check_credentials()
+
+ # Import here so missing kaggle package gives a clean error
+ try:
+ from kaggle.api.kaggle_api_extended import KaggleApiExtended
+ api = KaggleApiExtended()
+ api.authenticate()
+ except ImportError:
+ fatal("kaggle package not installed.\n Run: pip install kaggle")
+ except Exception as e:
+ fatal(f"Kaggle authentication failed: {e}\n"
+ "Make sure ~/.kaggle/kaggle.json exists and is valid.")
+
+ username = get_kaggle_username()
+ log(f"Authenticated as: {username}", "✓")
+
+ dataset_dir = Path(args.dataset_dir)
+ notebook_path = Path(args.notebook)
+ output_dir = Path(args.output_dir)
+
+ if not dataset_dir.exists():
+ fatal(f"Dataset directory not found: {dataset_dir}")
+ if not notebook_path.exists():
+ fatal(f"Notebook not found: {notebook_path}")
+
+ print(f"""
+╔══════════════════════════════════════════════════════╗
+║ Kaggle Training Automation ║
+╠══════════════════════════════════════════════════════╣
+ Kaggle user : {username}
+ Dataset : {args.dataset_name} ({dataset_dir})
+ Kernel : {args.kernel_name}
+ Notebook : {notebook_path}
+ GPU : {'yes (T4)' if gpu else 'no (CPU)'}
+ Output : {output_dir}
+ Poll interval : {args.poll_interval}s
+ Timeout : {args.timeout} min
+╚══════════════════════════════════════════════════════╝
+""")
+
+ with tempfile.TemporaryDirectory() as tmp_str:
+ tmp = Path(tmp_str)
+
+ # ── Step 1: upload dataset ────────────────────────────────────────────
+ if not args.skip_upload:
+ dataset_ref = upload_dataset(
+ api, username, args.dataset_name,
+ dataset_dir, tmp, args.public
+ )
+ # Give Kaggle a moment to process the upload before pushing the kernel
+ log("Waiting 10 s for dataset to register on Kaggle…", "⏳")
+ time.sleep(10)
+ else:
+ dataset_ref = f"{username}/{slugify(args.dataset_name)}"
+ log(f"Skipping upload — using existing dataset: {dataset_ref}", "↷")
+
+ # ── Step 2: patch notebook ────────────────────────────────────────────
+ patched_nb = patch_notebook(notebook_path, dataset_ref, tmp)
+
+ # ── Step 3: push kernel ───────────────────────────────────────────────
+ kernel_ref = f"{username}/{slugify(args.kernel_name)}"
+ if not args.skip_push:
+ kernel_ref = push_kernel(
+ api, username, args.kernel_name,
+ patched_nb, dataset_ref, tmp,
+ gpu, args.public,
+ )
+ else:
+ log(f"Skipping push — polling existing kernel: {kernel_ref}", "↷")
+
+ # ── Step 4: poll ──────────────────────────────────────────────────────
+ success = poll_kernel(api, kernel_ref, args.poll_interval, args.timeout)
+
+ if not success:
+ log(
+ f"Training did not complete successfully.\n"
+ f" View logs : https://www.kaggle.com/code/{kernel_ref}\n"
+ f" Re-run download only:\n"
+ f" python kaggle_train.py --skip_upload --skip_push "
+ f"--kernel_name {args.kernel_name}",
+ "✗"
+ )
+ sys.exit(1)
+
+ # ── Step 5: download ──────────────────────────────────────────────────
+ downloaded = download_outputs(api, kernel_ref, output_dir)
+
+ # ── Final summary ─────────────────────────────────────────────────────────
+ best_pt = next((f for f in downloaded if f.name == "best.pt"), None)
+ cls_json = next((f for f in downloaded if f.name == "class_names.json"), None)
+ curves_png = next((f for f in downloaded if "curves" in f.name), None)
+
+ print()
+ print("═" * 55)
+ print(" Training pipeline complete!")
+ if best_pt:
+ print(f" Weights : {best_pt}")
+ if cls_json:
+ print(f" Classes : {cls_json}")
+ if curves_png:
+ print(f" Curves : {curves_png}")
+ print()
+ print(" Next steps:")
+ print(f" python 3_inference.py \\")
+ print(f" --classifier_weights {output_dir}/best.pt \\")
+ print(f" --detector_weights detector/best.pt \\")
+ print(f" --source shelf_images/")
+ print("═" * 55)
+
+
+if __name__ == "__main__":
+ main()