Spaces:

Moni38
/

smartvision-ai

Sleeping

Priya Monisha

fix: pass ruff lint+format, clean README, remove stray script

fa0ce5b 5 months ago

11 kB

# filename: src/data/loader.py
# purpose:  HuggingFace COCO streaming, collection, and dataset saving pipeline
# version:  2.0

# SCHEMA NOTES (detection-datasets/coco):
#   sample["image"]                → PIL.Image (direct, not bytes dict)
#   sample["objects"]["bbox"]      → list of [x, y, w, h] (top-left + width/height, absolute pixels)
#   sample["objects"]["category"]  → list of 0-indexed HF category IDs (0-79, NOT COCO IDs 1-90)
#   No trust_remote_code needed. No HF token required for public dataset.
#   bbox format confirmed via mentor notebook: crop((x, y, x+w, y+h)) produces correct results.

import logging
from pathlib import Path

from PIL import Image

from config import (
    CLASSES,
    CHECKPOINT_FILE,
    HF_CATEGORY_TO_CLASS_IDX,
    NUM_CLASSES,
    SELECTED_CLASSES,
    TRAIN_SPLIT,
    VAL_SPLIT,
)
from src.utils.helpers import load_json, save_json

logger = logging.getLogger(__name__)


# ── COCO mapping verification ─────────────────────────────────────────────────
# Rule 28: call this as the first thing in 01_data_acquisition.py


def verify_coco_mapping() -> None:
    """Verify SELECTED_CLASSES and HF_CATEGORY_TO_CLASS_IDX are consistent."""
    assert len(SELECTED_CLASSES) == NUM_CLASSES, (
        f"SELECTED_CLASSES has {len(SELECTED_CLASSES)} entries, expected {NUM_CLASSES}"
    )
    assert len(HF_CATEGORY_TO_CLASS_IDX) == NUM_CLASSES, (
        f"HF_CATEGORY_TO_CLASS_IDX has {len(HF_CATEGORY_TO_CLASS_IDX)} entries, expected {NUM_CLASSES}"
    )
    # Check round-trip: class_name → hf_id → class_idx → class_name
    for cls_name, hf_id in SELECTED_CLASSES.items():
        assert hf_id in HF_CATEGORY_TO_CLASS_IDX, (
            f"HF ID {hf_id} (for '{cls_name}') missing from HF_CATEGORY_TO_CLASS_IDX"
        )
        class_idx = HF_CATEGORY_TO_CLASS_IDX[hf_id]
        assert CLASSES[class_idx] == cls_name, (
            f"Round-trip mismatch: '{cls_name}' → hf_id={hf_id} → idx={class_idx} → '{CLASSES[class_idx]}'"
        )
    # No duplicate HF IDs
    hf_ids = list(SELECTED_CLASSES.values())
    assert len(hf_ids) == len(set(hf_ids)), (
        "Duplicate HF category IDs in SELECTED_CLASSES"
    )
    logger.info(f"COCO mapping verified: {NUM_CLASSES} classes, no duplicates")
    print(f"COCO mapping OK: {NUM_CLASSES} classes, HF 0-indexed IDs, no duplicates")


# ── Checkpoint / resume ───────────────────────────────────────────────────────


def load_checkpoint() -> dict[str, int]:
    if CHECKPOINT_FILE.exists():
        data = load_json(CHECKPOINT_FILE)
        total = sum(data.values())
        logger.info(f"Resumed from checkpoint: {total} images already collected")
        return data
    return {cls: 0 for cls in CLASSES}


def save_checkpoint(progress: dict[str, int]) -> None:
    CHECKPOINT_FILE.parent.mkdir(parents=True, exist_ok=True)
    save_json(progress, CHECKPOINT_FILE)


# ── YOLO bbox conversion ──────────────────────────────────────────────────────
# bbox in HF dataset: [x, y, w, h] — top-left corner + width/height (absolute pixels)
# Confirmed by mentor notebook: crop((x, y, x+w, y+h)) produces correct object crops


def bbox_to_yolo(
    x: float,
    y: float,
    w: float,
    h: float,
    img_w: int,
    img_h: int,
    class_idx: int,
) -> list[float]:
    """Convert [x, y, w, h] absolute pixels → YOLO [class, x_center, y_center, w, h] normalized."""
    x_center = (x + w / 2.0) / img_w
    y_center = (y + h / 2.0) / img_h
    w_norm = w / img_w
    h_norm = h / img_h
    # Clamp to valid range
    x_center = max(0.0, min(1.0, x_center))
    y_center = max(0.0, min(1.0, y_center))
    w_norm = max(0.001, min(1.0, w_norm))
    h_norm = max(0.001, min(1.0, h_norm))
    return [class_idx, x_center, y_center, w_norm, h_norm]


# ── Overlap utilities ─────────────────────────────────────────────────────────


def compute_overlap_ratio(target_box: list, other_box: list) -> float:
    """
    Fraction of target_box area covered by other_box.
    Expects [x, y, w, h] format (top-left corner + dimensions).
    Use instead of IoU when checking whether a large object (person) dominates
    a smaller object (chair). IoU penalises the large union area and under-detects
    containment when other_box >> target_box.
    Returns float in [0, 1].
    """
    assert target_box[2] > 0 and target_box[3] > 0, (
        f"target_box width/height must be positive: {target_box}. "
        "Ensure bbox format is [x, y, w, h], not [x1, y1, x2, y2]."
    )
    assert other_box[2] > 0 and other_box[3] > 0, (
        f"other_box width/height must be positive: {other_box}. "
        "Ensure bbox format is [x, y, w, h], not [x1, y1, x2, y2]."
    )
    tx1, ty1 = target_box[0], target_box[1]
    tx2, ty2 = target_box[0] + target_box[2], target_box[1] + target_box[3]
    ox1, oy1 = other_box[0], other_box[1]
    ox2, oy2 = other_box[0] + other_box[2], other_box[1] + other_box[3]

    inter_x1 = max(tx1, ox1)
    inter_y1 = max(ty1, oy1)
    inter_x2 = min(tx2, ox2)
    inter_y2 = min(ty2, oy2)
    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)

    target_area = target_box[2] * target_box[3]
    return inter_area / target_area if target_area > 0 else 0.0


def check_person_overlap(
    target_bbox: list,
    all_bboxes: list,
    all_cat_ids: list,
    person_cat_id: int,
    threshold: float = 0.50,
) -> bool:
    """
    Return True if any person bbox covers > threshold fraction of target_bbox area.
    Uses compute_overlap_ratio, NOT IoU. IoU under-detects containment when
    person bbox is larger than chair bbox (the common COCO case).
    person_cat_id: must match the ID system in all_cat_ids.
    In this project: all_cat_ids contains HF 0-indexed IDs (loader.py line 8).
    Pass SELECTED_CLASSES["person"] = 0.
    threshold: 0.50 default (conservative). Adjust to 0.30 in a second
    re-collection run if chair accuracy remains low after first run.
    """
    for bbox, cat_id in zip(all_bboxes, all_cat_ids):
        if cat_id != person_cat_id:
            continue
        if compute_overlap_ratio(target_bbox, bbox) > threshold:
            return True
    return False


# ── Crop quality gates ────────────────────────────────────────────────────────


def check_crop_quality(img: Image.Image, x1: int, y1: int, x2: int, y2: int) -> str:
    """Return 'ok' or a rejection reason string.
    Gates (all three must pass):
      too_small  : crop < 48×48px in original image (upscale > 4.7× → blurry)
      too_tiny   : bbox covers < 1% of image area (invisible/background object)
      bad_aspect : longer side > 5× shorter side (sliver or letterbox crop)
    """
    crop_w = x2 - x1
    crop_h = y2 - y1
    if crop_w < 80 or crop_h < 80:
        return "too_small"
    if (crop_w * crop_h) / (img.width * img.height) < 0.02:
        return "too_tiny"
    if max(crop_w, crop_h) / min(crop_w, crop_h) > 5.0:
        return "bad_aspect"
    return "ok"


# ── Crop and save classification image ───────────────────────────────────────


def save_classification_crop(
    img: Image.Image,
    bbox: list[float],
    class_name: str,
    split: str,
    img_idx: int,
    classification_dir: Path,
    crop_size: int = 224,
) -> bool | str:
    """Crop object bbox, apply quality gates, resize to 224×224, save.
    Returns:
        True          — saved successfully
        "too_small"   — crop < 48px in either dimension
        "too_tiny"    — bbox covers < 1% of image area
        "bad_aspect"  — aspect ratio > 5:1
        False         — unexpected save error
    """
    x, y, w, h = bbox
    x1 = max(0, int(x))
    y1 = max(0, int(y))
    x2 = min(img.width, int(x + w))
    y2 = min(img.height, int(y + h))

    if x2 <= x1 or y2 <= y1:
        return "too_small"

    quality = check_crop_quality(img, x1, y1, x2, y2)
    if quality != "ok":
        return quality

    try:
        crop = img.crop((x1, y1, x2, y2))
        crop = crop.resize((crop_size, crop_size), Image.LANCZOS)
        out_dir = classification_dir / split / class_name
        out_dir.mkdir(parents=True, exist_ok=True)
        out_path = out_dir / f"{class_name}_{split}_{img_idx:04d}.jpg"
        crop.save(str(out_path), "JPEG", quality=95)
        return True
    except Exception as e:
        logger.warning(f"Crop failed for {class_name} idx={img_idx}: {e}")
        return False


# ── Save detection sample ─────────────────────────────────────────────────────


def save_detection_sample(
    img: Image.Image,
    annotations: dict,
    img_id: int,
    split: str,
    detection_dir: Path,
) -> None:
    """Save full image and YOLO-format .txt label for detection task."""
    img_w, img_h = img.size
    bboxes = annotations["bbox"]
    cat_ids = annotations["category"]

    # Build YOLO annotation lines for our 25 classes only
    yolo_lines = []
    for bbox, cat_id in zip(bboxes, cat_ids):
        if cat_id not in HF_CATEGORY_TO_CLASS_IDX:
            continue
        class_idx = HF_CATEGORY_TO_CLASS_IDX[cat_id]
        x, y, w, h = bbox
        vals = bbox_to_yolo(x, y, w, h, img_w, img_h, class_idx)
        yolo_lines.append(
            f"{int(vals[0])} {vals[1]:.6f} {vals[2]:.6f} {vals[3]:.6f} {vals[4]:.6f}"
        )

    if not yolo_lines:
        return

    img_dir = detection_dir / "images" / split
    lbl_dir = detection_dir / "labels" / split
    img_dir.mkdir(parents=True, exist_ok=True)
    lbl_dir.mkdir(parents=True, exist_ok=True)

    img.save(str(img_dir / f"image_{img_id:06d}.jpg"), "JPEG", quality=95)
    (lbl_dir / f"image_{img_id:06d}.txt").write_text(
        "\n".join(yolo_lines), encoding="utf-8"
    )


# ── Split helper ──────────────────────────────────────────────────────────────


def get_split(img_idx: int, total: int) -> str:
    """Return 'train', 'val', or 'test' based on img_idx position."""
    train_end = int(total * TRAIN_SPLIT)
    val_end = train_end + int(total * VAL_SPLIT)
    if img_idx < train_end:
        return "train"
    elif img_idx < val_end:
        return "val"
    return "test"

	# filename: src/data/loader.py
	# purpose: HuggingFace COCO streaming, collection, and dataset saving pipeline
	# version: 2.0

	# SCHEMA NOTES (detection-datasets/coco):
	# sample["image"] → PIL.Image (direct, not bytes dict)
	# sample["objects"]["bbox"] → list of [x, y, w, h] (top-left + width/height, absolute pixels)
	# sample["objects"]["category"] → list of 0-indexed HF category IDs (0-79, NOT COCO IDs 1-90)
	# No trust_remote_code needed. No HF token required for public dataset.
	# bbox format confirmed via mentor notebook: crop((x, y, x+w, y+h)) produces correct results.

	import logging
	from pathlib import Path

	from PIL import Image

	from config import (
	CLASSES,
	CHECKPOINT_FILE,
	HF_CATEGORY_TO_CLASS_IDX,
	NUM_CLASSES,
	SELECTED_CLASSES,
	TRAIN_SPLIT,
	VAL_SPLIT,
	)
	from src.utils.helpers import load_json, save_json

	logger = logging.getLogger(__name__)


	# ── COCO mapping verification ─────────────────────────────────────────────────
	# Rule 28: call this as the first thing in 01_data_acquisition.py


	def verify_coco_mapping() -> None:
	"""Verify SELECTED_CLASSES and HF_CATEGORY_TO_CLASS_IDX are consistent."""
	assert len(SELECTED_CLASSES) == NUM_CLASSES, (
	f"SELECTED_CLASSES has {len(SELECTED_CLASSES)} entries, expected {NUM_CLASSES}"
	)
	assert len(HF_CATEGORY_TO_CLASS_IDX) == NUM_CLASSES, (
	f"HF_CATEGORY_TO_CLASS_IDX has {len(HF_CATEGORY_TO_CLASS_IDX)} entries, expected {NUM_CLASSES}"
	)
	# Check round-trip: class_name → hf_id → class_idx → class_name
	for cls_name, hf_id in SELECTED_CLASSES.items():
	assert hf_id in HF_CATEGORY_TO_CLASS_IDX, (
	f"HF ID {hf_id} (for '{cls_name}') missing from HF_CATEGORY_TO_CLASS_IDX"
	)
	class_idx = HF_CATEGORY_TO_CLASS_IDX[hf_id]
	assert CLASSES[class_idx] == cls_name, (
	f"Round-trip mismatch: '{cls_name}' → hf_id={hf_id} → idx={class_idx} → '{CLASSES[class_idx]}'"
	)
	# No duplicate HF IDs
	hf_ids = list(SELECTED_CLASSES.values())
	assert len(hf_ids) == len(set(hf_ids)), (
	"Duplicate HF category IDs in SELECTED_CLASSES"
	)
	logger.info(f"COCO mapping verified: {NUM_CLASSES} classes, no duplicates")
	print(f"COCO mapping OK: {NUM_CLASSES} classes, HF 0-indexed IDs, no duplicates")


	# ── Checkpoint / resume ───────────────────────────────────────────────────────


	def load_checkpoint() -> dict[str, int]:
	if CHECKPOINT_FILE.exists():
	data = load_json(CHECKPOINT_FILE)
	total = sum(data.values())
	logger.info(f"Resumed from checkpoint: {total} images already collected")
	return data
	return {cls: 0 for cls in CLASSES}


	def save_checkpoint(progress: dict[str, int]) -> None:
	CHECKPOINT_FILE.parent.mkdir(parents=True, exist_ok=True)
	save_json(progress, CHECKPOINT_FILE)


	# ── YOLO bbox conversion ──────────────────────────────────────────────────────
	# bbox in HF dataset: [x, y, w, h] — top-left corner + width/height (absolute pixels)
	# Confirmed by mentor notebook: crop((x, y, x+w, y+h)) produces correct object crops


	def bbox_to_yolo(
	x: float,
	y: float,
	w: float,
	h: float,
	img_w: int,
	img_h: int,
	class_idx: int,
	) -> list[float]:
	"""Convert [x, y, w, h] absolute pixels → YOLO [class, x_center, y_center, w, h] normalized."""
	x_center = (x + w / 2.0) / img_w
	y_center = (y + h / 2.0) / img_h
	w_norm = w / img_w
	h_norm = h / img_h
	# Clamp to valid range
	x_center = max(0.0, min(1.0, x_center))
	y_center = max(0.0, min(1.0, y_center))
	w_norm = max(0.001, min(1.0, w_norm))
	h_norm = max(0.001, min(1.0, h_norm))
	return [class_idx, x_center, y_center, w_norm, h_norm]


	# ── Overlap utilities ─────────────────────────────────────────────────────────


	def compute_overlap_ratio(target_box: list, other_box: list) -> float:
	"""
	Fraction of target_box area covered by other_box.
	Expects [x, y, w, h] format (top-left corner + dimensions).

	Use instead of IoU when checking whether a large object (person) dominates
	a smaller object (chair). IoU penalises the large union area and under-detects
	containment when other_box >> target_box.

	Returns float in [0, 1].
	"""
	assert target_box[2] > 0 and target_box[3] > 0, (
	f"target_box width/height must be positive: {target_box}. "
	"Ensure bbox format is [x, y, w, h], not [x1, y1, x2, y2]."
	)
	assert other_box[2] > 0 and other_box[3] > 0, (
	f"other_box width/height must be positive: {other_box}. "
	"Ensure bbox format is [x, y, w, h], not [x1, y1, x2, y2]."
	)
	tx1, ty1 = target_box[0], target_box[1]
	tx2, ty2 = target_box[0] + target_box[2], target_box[1] + target_box[3]
	ox1, oy1 = other_box[0], other_box[1]
	ox2, oy2 = other_box[0] + other_box[2], other_box[1] + other_box[3]

	inter_x1 = max(tx1, ox1)
	inter_y1 = max(ty1, oy1)
	inter_x2 = min(tx2, ox2)
	inter_y2 = min(ty2, oy2)
	inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)

	target_area = target_box[2] * target_box[3]
	return inter_area / target_area if target_area > 0 else 0.0


	def check_person_overlap(
	target_bbox: list,
	all_bboxes: list,
	all_cat_ids: list,
	person_cat_id: int,
	threshold: float = 0.50,
	) -> bool:
	"""
	Return True if any person bbox covers > threshold fraction of target_bbox area.

	Uses compute_overlap_ratio, NOT IoU. IoU under-detects containment when
	person bbox is larger than chair bbox (the common COCO case).

	person_cat_id: must match the ID system in all_cat_ids.
	In this project: all_cat_ids contains HF 0-indexed IDs (loader.py line 8).
	Pass SELECTED_CLASSES["person"] = 0.

	threshold: 0.50 default (conservative). Adjust to 0.30 in a second
	re-collection run if chair accuracy remains low after first run.
	"""
	for bbox, cat_id in zip(all_bboxes, all_cat_ids):
	if cat_id != person_cat_id:
	continue
	if compute_overlap_ratio(target_bbox, bbox) > threshold:
	return True
	return False


	# ── Crop quality gates ────────────────────────────────────────────────────────


	def check_crop_quality(img: Image.Image, x1: int, y1: int, x2: int, y2: int) -> str:
	"""Return 'ok' or a rejection reason string.

	Gates (all three must pass):
	too_small : crop < 48×48px in original image (upscale > 4.7× → blurry)
	too_tiny : bbox covers < 1% of image area (invisible/background object)
	bad_aspect : longer side > 5× shorter side (sliver or letterbox crop)
	"""
	crop_w = x2 - x1
	crop_h = y2 - y1
	if crop_w < 80 or crop_h < 80:
	return "too_small"
	if (crop_w * crop_h) / (img.width * img.height) < 0.02:
	return "too_tiny"
	if max(crop_w, crop_h) / min(crop_w, crop_h) > 5.0:
	return "bad_aspect"
	return "ok"


	# ── Crop and save classification image ───────────────────────────────────────


	def save_classification_crop(
	img: Image.Image,
	bbox: list[float],
	class_name: str,
	split: str,
	img_idx: int,
	classification_dir: Path,
	crop_size: int = 224,
	) -> bool \| str:
	"""Crop object bbox, apply quality gates, resize to 224×224, save.

	Returns:
	True — saved successfully
	"too_small" — crop < 48px in either dimension
	"too_tiny" — bbox covers < 1% of image area
	"bad_aspect" — aspect ratio > 5:1
	False — unexpected save error
	"""
	x, y, w, h = bbox
	x1 = max(0, int(x))
	y1 = max(0, int(y))
	x2 = min(img.width, int(x + w))
	y2 = min(img.height, int(y + h))

	if x2 <= x1 or y2 <= y1:
	return "too_small"

	quality = check_crop_quality(img, x1, y1, x2, y2)
	if quality != "ok":
	return quality

	try:
	crop = img.crop((x1, y1, x2, y2))
	crop = crop.resize((crop_size, crop_size), Image.LANCZOS)
	out_dir = classification_dir / split / class_name
	out_dir.mkdir(parents=True, exist_ok=True)
	out_path = out_dir / f"{class_name}_{split}_{img_idx:04d}.jpg"
	crop.save(str(out_path), "JPEG", quality=95)
	return True
	except Exception as e:
	logger.warning(f"Crop failed for {class_name} idx={img_idx}: {e}")
	return False


	# ── Save detection sample ─────────────────────────────────────────────────────


	def save_detection_sample(
	img: Image.Image,
	annotations: dict,
	img_id: int,
	split: str,
	detection_dir: Path,
	) -> None:
	"""Save full image and YOLO-format .txt label for detection task."""
	img_w, img_h = img.size
	bboxes = annotations["bbox"]
	cat_ids = annotations["category"]

	# Build YOLO annotation lines for our 25 classes only
	yolo_lines = []
	for bbox, cat_id in zip(bboxes, cat_ids):
	if cat_id not in HF_CATEGORY_TO_CLASS_IDX:
	continue
	class_idx = HF_CATEGORY_TO_CLASS_IDX[cat_id]
	x, y, w, h = bbox
	vals = bbox_to_yolo(x, y, w, h, img_w, img_h, class_idx)
	yolo_lines.append(
	f"{int(vals[0])} {vals[1]:.6f} {vals[2]:.6f} {vals[3]:.6f} {vals[4]:.6f}"
	)

	if not yolo_lines:
	return

	img_dir = detection_dir / "images" / split
	lbl_dir = detection_dir / "labels" / split
	img_dir.mkdir(parents=True, exist_ok=True)
	lbl_dir.mkdir(parents=True, exist_ok=True)

	img.save(str(img_dir / f"image_{img_id:06d}.jpg"), "JPEG", quality=95)
	(lbl_dir / f"image_{img_id:06d}.txt").write_text(
	"\n".join(yolo_lines), encoding="utf-8"
	)


	# ── Split helper ──────────────────────────────────────────────────────────────


	def get_split(img_idx: int, total: int) -> str:
	"""Return 'train', 'val', or 'test' based on img_idx position."""
	train_end = int(total * TRAIN_SPLIT)
	val_end = train_end + int(total * VAL_SPLIT)
	if img_idx < train_end:
	return "train"
	elif img_idx < val_end:
	return "val"
	return "test"