| """Validate examples, group duplicates, and load target-specific labels.""" |
|
|
| import hashlib |
| import json |
| import os |
| import random |
| import re |
| from collections import Counter, defaultdict |
| from functools import partial |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import torch |
| from PIL import Image |
| from torch.utils.data import Dataset |
| from torchvision import transforms as T |
|
|
| from targets import ALIASES, BACKGROUND, target_name |
|
|
|
|
| def read_rgb(path): |
| with Image.open(path) as image: |
| return image.convert("RGB") |
|
|
|
|
| def pixel_hash(path): |
| image = read_rgb(path) |
| return hashlib.sha256(str(image.size).encode() + image.tobytes()).hexdigest() |
|
|
|
|
| def average_hash(path): |
| """Survives rescaling and pixel noise, so noisy copies of a tile group together.""" |
| with Image.open(path) as image: |
| return _average_hash(image) |
|
|
|
|
| def _average_hash(image): |
| pixels = image.convert("L").resize((8, 8), Image.Resampling.BOX).tobytes() |
| mean = sum(pixels) / 64 |
| return "a" + format(int("".join("1" if v > mean else "0" for v in pixels), 2), "016x") |
|
|
|
|
| GRID_SIZES = {"3x3": 9, "4x4": 16} |
| IMAGE_NAMES = ("image.jpg", "image.jpeg", "image.png", "image.webp") |
|
|
|
|
| def grid_boxes(width, height, tiles): |
| """Cell boxes, left to right then top to bottom, with edges rounded like the labeling tools.""" |
| n = int(tiles**0.5) |
| xs = [round(width * i / n) for i in range(n + 1)] |
| ys = [round(height * i / n) for i in range(n + 1)] |
| return [(xs[c], ys[r], xs[c + 1], ys[r + 1]) for r in range(n) for c in range(n)] |
|
|
|
|
| def image_hash(image): |
| return hashlib.sha256(str(image.size).encode() + image.tobytes()).hexdigest() |
|
|
|
|
| def record_image_path(root, record): |
| """Locate a full image from either a published dataset or a legacy manifest.""" |
| if "image_path" in record: |
| return Path(root) / record["image_path"] |
| return Path(root) / record["id"] / record.get("image", "full_images.png") |
|
|
|
|
| def prepare_jsonl(root, output): |
| """Import the published dataset's existing splits without reshuffling them.""" |
| root = Path(root).resolve() |
| data = root / "data" if (root / "data" / "3x3").is_dir() else root |
| records, counts = [], Counter() |
| for grid, tiles in GRID_SIZES.items(): |
| for split in ("train", "validation", "test", "review", "excluded"): |
| metadata = data / grid / split / "metadata.jsonl" |
| if not metadata.is_file(): |
| raise ValueError(f"Missing dataset metadata: {metadata}") |
| with metadata.open() as source: |
| for line_number, line in enumerate(source, 1): |
| try: |
| row = json.loads(line) |
| if row["grid"] != grid or not isinstance(row["target"], str): |
| raise ValueError("invalid grid or target") |
| answers, uncertain = row["correct_answers"], row["uncertain_cells"] |
| if any(not isinstance(x, list) for x in (answers, uncertain)) or any( |
| type(i) is not int or not 0 <= i < tiles for i in answers + uncertain |
| ): |
| raise ValueError("invalid cell indices") |
| path = (metadata.parent / row["file_name"]).resolve() |
| if not path.is_relative_to(metadata.parent.resolve()) or not path.is_file(): |
| raise ValueError("invalid or missing image") |
| counts[split] += 1 |
| if split == "excluded": |
| continue |
| record = { |
| "id": row["id"], "tiles": tiles, |
| "target": target_name(row["target"]), |
| "answers": sorted(set(answers)), "uncertain": sorted(set(uncertain)), |
| "image_path": str(path.relative_to(data)), "tile_files": False, |
| "label_source": row.get("label_source"), |
| "split": "val" if split == "validation" else split, |
| } |
| if split == "review" and isinstance(row.get("machine_answers"), list): |
| record["labeler_answers"] = row["machine_answers"] |
| records.append(record) |
| except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: |
| raise ValueError(f"Invalid record in {metadata}:{line_number}: {exc}") from exc |
| if not records: |
| raise ValueError(f"No labeled examples in {data}") |
| manifest = {"root": str(data), "records": records} |
| output = Path(output) |
| output.mkdir(parents=True, exist_ok=True) |
| (output / "split.json").write_text(json.dumps(manifest, indent=2) + "\n") |
| summary = {"splits": dict(counts), "uncertain_cells": sum(len(r["uncertain"]) for r in records)} |
| (output / "data_summary.json").write_text(json.dumps(summary, indent=2) + "\n") |
| return manifest, summary |
|
|
|
|
| def _scan_labeled(directory, sam=False): |
| """One <item>/{image.jpg, metadata.json} folder: (record, None) or (None, reason). |
| With sam, 4x4 answers come from <item>/sam.json instead.""" |
| try: |
| info = json.loads((directory / "metadata.json").read_text()) |
| except (OSError, json.JSONDecodeError): |
| return None, "unreadable metadata.json" |
| if not isinstance(info, dict): |
| return None, "invalid metadata.json" |
| if not isinstance(info.get("target"), str) or not info["target"].strip(): |
| return None, "invalid target" |
| n = GRID_SIZES.get(str(info.get("grid", "")).lower()) |
| if n is None: |
| return None, "unknown grid" |
| if sam and n == 16: |
| try: |
| label = json.loads((directory / "sam.json").read_text()) |
| except (OSError, json.JSONDecodeError): |
| return None, "no sam.json" |
| if not isinstance(label, dict) or not {"correct_answers", "model"} <= label.keys(): |
| return None, "invalid sam.json" |
| |
| info = { |
| **info, |
| "correct_answers": label["correct_answers"], |
| "uncertain_cells": [], |
| "label_source": f"{label['model']}:sam", |
| } |
| if "correct_answers" not in info: |
| return None, "not labeled" |
| answers, uncertain = info["correct_answers"], info.get("uncertain_cells", []) |
| if ( |
| not isinstance(answers, list) |
| or not isinstance(uncertain, list) |
| or any(type(i) is not int or not 0 <= i < n for i in answers + uncertain) |
| ): |
| return None, "invalid answers" |
| path = next((directory / name for name in IMAGE_NAMES if (directory / name).exists()), None) |
| if path is None: |
| return None, "no image" |
| try: |
| image = read_rgb(path) |
| except OSError: |
| return None, "unreadable image" |
| hashes = [image_hash(image)] |
| if n == 9: |
| |
| hashes += [image_hash(image.crop(box)) for box in grid_boxes(*image.size, n)] |
| return { |
| "id": directory.name, |
| "tiles": n, |
| "target": target_name(info["target"]), |
| "answers": sorted(set(answers)), |
| "uncertain": sorted(set(uncertain)), |
| "image": path.name, |
| "tile_files": False, |
| "label_source": info.get("label_source"), |
| "hashes": hashes, |
| }, None |
|
|
|
|
| def scan(root, workers=8, sam=False): |
| """Answer-key folders (info.json + full_images.png + tile_N.png) and labeled |
| <item>/{image.jpg, metadata.json} folders. Returns (records, {skip reason: count}).""" |
| if workers < 1: |
| raise ValueError("Scanning workers must be positive") |
| directories = sorted( |
| p for p in Path(root).iterdir() if p.is_dir() and not p.name.startswith(".") |
| ) |
| labeled = [d for d in directories if not (d / "info.json").exists()] |
| records, skipped = [], Counter() |
| if labeled: |
| with Pool(workers) as pool: |
| for record, reason in pool.imap(partial(_scan_labeled, sam=sam), labeled, chunksize=64): |
| if record: |
| records.append(record) |
| else: |
| skipped[reason] += 1 |
| for directory in directories: |
| if not (directory / "info.json").exists(): |
| continue |
| info = json.loads((directory / "info.json").read_text()) |
| indices = sorted(int(p.stem.split("_")[1]) for p in directory.glob("tile_*.png")) |
| n = len(indices) |
| if n not in (9, 16) or indices != list(range(n)): |
| raise ValueError(f"Invalid tiles: {directory}") |
| answers = info["correct_answers"] |
| if ( |
| not isinstance(answers, list) |
| or any(type(i) is not int or not 0 <= i < n for i in answers) |
| or len(set(answers)) != len(answers) |
| ): |
| raise ValueError(f"Invalid answers: {directory}") |
| paths = [directory / "full_images.png"] + [directory / f"tile_{i}.png" for i in range(n)] |
| records.append( |
| { |
| "id": directory.name, |
| "tiles": n, |
| "target": target_name(info["keyword"]), |
| "answers": sorted(answers), |
| "hashes": [pixel_hash(p) for p in paths], |
| } |
| ) |
| if not records: |
| raise ValueError(f"No labeled examples in {root} ({dict(skipped)})") |
| records.sort(key=lambda r: r["id"]) |
| return records, dict(skipped) |
|
|
|
|
| def split_records(records, seed, holdouts=("val", "test"), review=()): |
| """Records whose id is in review go to a "review" split; unreviewed records sharing an image |
| or tile with them get split "dropped" so nothing they show leaks into training.""" |
| |
| review = set(review) |
| parent = list(range(len(records))) |
|
|
| def find(i): |
| while parent[i] != i: |
| parent[i] = parent[parent[i]] |
| i = parent[i] |
| return i |
|
|
| seen = {} |
| for i, record in enumerate(records): |
| for digest in record["hashes"]: |
| if digest in seen: |
| parent[find(i)] = find(seen[digest]) |
| seen[digest] = i |
| groups = defaultdict(list) |
| for i, record in enumerate(records): |
| groups[find(i)].append(record) |
| strata = defaultdict(list) |
| for group in groups.values(): |
| if review and any(r.get("id") in review for r in group): |
| for record in group: |
| record["split"] = "review" if record.get("id") in review else "dropped" |
| continue |
| strata[tuple(sorted({(r["tiles"], r["target"]) for r in group}))].append(group) |
| rng = random.Random(seed) |
| for key in sorted(strata): |
| items = strata[key] |
| rng.shuffle(items) |
| |
| holdout = max(1, round(len(items) * 0.1)) if len(items) >= len(holdouts) + 1 else 0 |
| for i, group in enumerate(items): |
| split = holdouts[i // holdout] if i < holdout * len(holdouts) else "train" |
| for record in group: |
| record["split"] = split |
| return records |
|
|
|
|
| def prepare(root, output, seed=42, workers=8, sam=False, review_keys=()): |
| """review_keys: reference answer files ({"items": {id: {"answers": [...]}}}). Those grids |
| get the reviewed answers and their own "review" split, never used for training or selection.""" |
| output = Path(output) |
| output.mkdir(parents=True, exist_ok=True) |
| records, skipped = scan(root, workers, sam) |
| review = {} |
| for path in review_keys: |
| review.update(json.loads(Path(path).read_text())["items"]) |
| for record in records: |
| if record["id"] in review: |
| |
| record.update( |
| labeler_answers=record["answers"], |
| answers=review[record["id"]]["answers"], |
| uncertain=[], |
| label_source="review", |
| ) |
| |
| duplicates = defaultdict(list) |
| for record in records: |
| if record["id"] in review: |
| continue |
| key = (record["hashes"][0], record["target"], record["tiles"]) |
| duplicates[key].append(record) |
| excluded = [ |
| r |
| for group in duplicates.values() |
| if len({tuple(r["answers"]) for r in group}) > 1 |
| for r in group |
| ] |
| excluded_ids = {r["id"] for r in excluded} |
| records = split_records( |
| [r for r in records if r["id"] not in excluded_ids], seed, review=review |
| ) |
| dropped = sum(r["split"] == "dropped" for r in records) |
| records = [r for r in records if r["split"] != "dropped"] |
| manifest = { |
| "seed": seed, |
| "root": str(Path(root).resolve()), |
| "records": records, |
| "excluded_conflicting_labels": excluded, |
| } |
| (output / "label_conflicts.json").write_text(json.dumps(excluded, indent=2) + "\n") |
| (output / "split.json").write_text(json.dumps(manifest, indent=2) + "\n") |
| counts = {} |
| for n in (9, 16): |
| counts[str(n)] = { |
| target: dict( |
| Counter(r["split"] for r in records if r["tiles"] == n and r["target"] == target) |
| ) |
| for target in sorted({r["target"] for r in records if r["tiles"] == n}) |
| } |
| counts["skipped"] = skipped |
| counts["excluded_conflicting_labels"] = len(excluded) |
| counts["dropped_sharing_images_with_review"] = dropped |
| counts["uncertain_cells"] = sum(len(r.get("uncertain", [])) for r in records) |
| (output / "data_summary.json").write_text(json.dumps(counts, indent=2) + "\n") |
| return manifest, counts |
|
|
|
|
| def _file_hashes(path): |
| with Image.open(path) as image: |
| width, height = image.size |
| thumbnail = image.convert("L").resize((32, 32), Image.Resampling.BOX).tobytes() |
| return str(path), pixel_hash(path), average_hash(path), max(width, height), thumbnail |
|
|
|
|
| def _held_out_hashes(item): |
| root, record = item |
| directory = root / record["id"] |
| image = read_rgb(record_image_path(root, record)) |
| pictures = [image] |
| if record.get("tile_files", True): |
| pictures.extend(read_rgb(directory / f"tile_{i}.png") for i in range(record["tiles"])) |
| else: |
| pictures.extend(image.crop(box) for box in grid_boxes(*image.size, record["tiles"])) |
| return [ |
| ( |
| _average_hash(picture), |
| picture.convert("L").resize((32, 32), Image.Resampling.BOX).tobytes(), |
| ) |
| for picture in pictures |
| ] |
|
|
|
|
| def same_picture(a, b, tolerance=12): |
| """Average hashes collide for similar layouts; confirm with a 32x32 mean absolute difference.""" |
| return sum(abs(x - y) for x, y in zip(a, b)) / len(a) < tolerance |
|
|
|
|
| IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".webp") |
| |
| |
| BOX = re.compile(r"(\d+)\s*(\d\.\d{6})\s+(\d\.\d{6})\s+(\d\.\d{6})\s+(\d\.\d{6})") |
|
|
|
|
| def label_file(path): |
| """YOLO layout: <root>/images/<Class>/x.png pairs with <root>/labels/<Class>/x.txt.""" |
| parts = path.parts |
| if "images" not in parts: |
| return None |
| i = len(parts) - 1 - parts[::-1].index("images") |
| candidate = Path(*parts[:i], "labels", *parts[i + 1 :]).with_suffix(".txt") |
| return candidate if candidate.exists() else None |
|
|
|
|
| def box_classes(path): |
| return sorted({int(match[0]) for match in BOX.findall(path.read_text())}) |
|
|
|
|
| def scan_folders(roots): |
| """Images labeled by their parent folder; <split>/ and images/ levels are ignored. |
| |
| YOLO class ids are mapped to the folder where each id appears most often, so boxes |
| for other classes add those targets to an image. |
| """ |
| paths, labels, boxes = [], [], [] |
| for root in roots: |
| for path in sorted(Path(root).rglob("*")): |
| if path.suffix.lower() in IMAGE_SUFFIXES and not path.name.startswith("."): |
| paths.append(path) |
| labels.append(target_name(path.parent.name)) |
| label = label_file(path) |
| boxes.append(box_classes(label) if label else []) |
| votes = defaultdict(Counter) |
| for label, ids in zip(labels, boxes): |
| for i in ids: |
| votes[i][label] += 1 |
| id_names = {i: counts.most_common(1)[0][0] for i, counts in votes.items()} |
| targets = [sorted({label, *(id_names[i] for i in ids)}) for label, ids in zip(labels, boxes)] |
| return paths, targets, {str(i): name for i, name in sorted(id_names.items())} |
|
|
|
|
| def prepare_folders(roots, old_manifest, output, seed=42, workers=8): |
| """Merge exact copies, group near-copies, and re-split one or more class-folder datasets. |
| |
| The folders' own train/val divisions are ignored because many images appear in both. |
| Images resembling held-out grids of the old dataset are dropped. |
| """ |
| roots = [Path(r).resolve() for r in ([roots] if isinstance(roots, (str, Path)) else roots)] |
| base = Path(os.path.commonpath(roots)) if len(roots) > 1 else roots[0] |
| output = Path(output) |
| output.mkdir(parents=True, exist_ok=True) |
| paths, path_targets, id_names = scan_folders(roots) |
| old = json.loads(Path(old_manifest).read_text()) |
| old_root = Path(old["root"]) |
| held_out = [(old_root, record) for record in old["records"] if record["split"] != "train"] |
| with Pool(workers) as pool: |
| hashed = pool.map(_file_hashes, paths, chunksize=64) |
| held_out_hashes = defaultdict(list) |
| for pictures in pool.map(_held_out_hashes, held_out, chunksize=64): |
| for ahash, thumbnail in pictures: |
| held_out_hashes[ahash].append(thumbnail) |
| by_pixels, thumbnails = {}, {} |
| sources = Counter() |
| for path, labels, (_, digest, ahash, side, thumbnail) in zip(paths, path_targets, hashed): |
| thumbnails[digest] = thumbnail |
| relative = path.relative_to(base) |
| if digest not in by_pixels: |
| sources[relative.parts[0] if len(roots) > 1 else roots[0].name] += 1 |
| record = by_pixels.setdefault( |
| digest, |
| { |
| "path": str(relative), |
| "targets": [], |
| "kind": "full" if side >= 400 else "tile", |
| "tiles": "full" if side >= 400 else "tile", |
| "hashes": [digest, ahash], |
| }, |
| ) |
| record["targets"].extend(t for t in labels if t not in record["targets"]) |
| records, leaked, conflicts = [], [], [] |
| for record in by_pixels.values(): |
| record["targets"].sort() |
| if any( |
| same_picture(thumbnails[record["hashes"][0]], other) |
| for other in held_out_hashes.get(record["hashes"][1], []) |
| ): |
| leaked.append(record["path"]) |
| else: |
| if BACKGROUND in record["targets"] and len(record["targets"]) > 1: |
| |
| conflicts.append(record["path"]) |
| record["targets"].remove(BACKGROUND) |
| record["target"] = record["targets"][0] |
| records.append(record) |
| if not records: |
| raise ValueError("No folder images remain after excluding held-out images") |
| split_records(records, seed, holdouts=("val",)) |
| for record in records: |
| del record["target"], record["tiles"] |
| summary = { |
| "roots": [str(r) for r in roots], |
| "files": len(paths), |
| "unique_images": len(by_pixels), |
| "unique_images_by_first_source": dict(sources), |
| "yolo_class_ids": id_names, |
| "exact_duplicates_merged": len(paths) - len(by_pixels), |
| "multi_label_images": sum(len(r["targets"]) > 1 for r in records), |
| "dropped_matching_old_heldout": len(leaked), |
| "background_label_overridden_by_target": len(conflicts), |
| "largest_near_duplicate_group": max(Counter(r["hashes"][1] for r in records).values()), |
| "counts": { |
| kind: { |
| target: dict( |
| Counter( |
| r["split"] for r in records if r["kind"] == kind and target in r["targets"] |
| ) |
| ) |
| for target in sorted({t for r in records for t in r["targets"]}) |
| } |
| for kind in ("tile", "full") |
| }, |
| } |
| manifest = { |
| "seed": seed, |
| "root": str(base), |
| "sources": [str(r) for r in roots], |
| "old_manifest": str(Path(old_manifest).resolve()), |
| "records": records, |
| "dropped_matching_old_heldout": leaked, |
| "background_label_overridden_by_target": conflicts, |
| } |
| (output / "folders.json").write_text(json.dumps(manifest, indent=2) + "\n") |
| (output / "folders_summary.json").write_text(json.dumps(summary, indent=2) + "\n") |
| return manifest, summary |
|
|
|
|
| class AddNoise: |
| """Some real tiles carry heavy pixel noise.""" |
|
|
| def __init__(self, probability=0.3, max_std=0.15): |
| self.probability, self.max_std = probability, max_std |
|
|
| def __call__(self, image): |
| if random.random() >= self.probability: |
| return image |
| return (image + torch.randn_like(image) * random.uniform(0.02, self.max_std)).clamp(0, 1) |
|
|
|
|
| def transform(size, training=False, noise=False): |
| operations = [T.Resize((size, size), interpolation=T.InterpolationMode.BILINEAR)] |
| if training: |
| operations.append(T.ColorJitter(brightness=0.15, contrast=0.15, saturation=0.1)) |
| operations.append(T.ToTensor()) |
| if training and noise: |
| operations.append(AddNoise()) |
| return T.Compose(operations + [T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) |
|
|
|
|
| class GridDataset(Dataset): |
| def __init__( |
| self, |
| root, |
| records, |
| classes, |
| size, |
| training=False, |
| tiles=False, |
| tile_grids=False, |
| include_full_images=False, |
| ): |
| self.root = Path(root) |
| self.records = records |
| self.classes = classes |
| self.transform = transform(size, training) |
| self.tiles = tiles |
| self.tile_grids = tile_grids |
| self.items = [(r, i) for r in records for i in (range(r["tiles"]) if tiles else [None])] |
| if include_full_images: |
| if not tiles: |
| raise ValueError("Whole-image classification examples require tile training") |
| self.items += [(r, -1) for r in records if r["tiles"] == 16] |
|
|
| def __len__(self): |
| return len(self.items) |
|
|
| def full_image(self, record): |
| return read_rgb(record_image_path(self.root, record)) |
|
|
| def tile_images(self, record, indices): |
| if record.get("tile_files", True): |
| return [read_rgb(self.root / record["id"] / f"tile_{i}.png") for i in indices] |
| image = self.full_image(record) |
| boxes = grid_boxes(*image.size, record["tiles"]) |
| return [image.crop(boxes[i]) for i in indices] |
|
|
| def __getitem__(self, index): |
| record, tile = self.items[index] |
| label = torch.tensor([float(i in record["answers"]) for i in range(record["tiles"])]) |
| if tile == -1: |
| image = self.transform(self.full_image(record)) |
| label = label.any().float() |
| elif tile is not None: |
| image = self.transform(self.tile_images(record, [tile])[0]) |
| label = label[tile] |
| elif record["tiles"] == 9 or self.tile_grids: |
| image = torch.stack( |
| [self.transform(t) for t in self.tile_images(record, range(record["tiles"]))] |
| ) |
| else: |
| image = self.transform(self.full_image(record)) |
| return image, self.classes.index(record["target"]), label |
|
|
|
|
| def supervision(classes, cells, target, cell_labels=None, any_label=None): |
| """Labels and weights for (classes, cells) logits; zero weight means unknown.""" |
| labels, weights = torch.zeros(len(classes), cells), torch.zeros(len(classes), cells) |
| any_labels, any_weights = torch.zeros(len(classes)), torch.zeros(len(classes)) |
| if cell_labels is not None: |
| labels[target], weights[target] = cell_labels, 1 |
| if any_label is not None: |
| any_labels[target], any_weights[target] = any_label, 1 |
| return labels, weights, any_labels, any_weights |
|
|
|
|
| class SupervisedGrids(Dataset): |
| """Answer-key items in the per-class format: only the queried target is known. |
| Cells listed as uncertain get zero weight unless keep_uncertain is set.""" |
|
|
| def __init__(self, grids, cells, keep_uncertain=False, flip=False): |
| self.grids, self.cells, self.keep_uncertain, self.flip = grids, cells, keep_uncertain, flip |
|
|
| def __len__(self): |
| return len(self.grids) |
|
|
| def __getitem__(self, index): |
| record, tile = self.grids.items[index] |
| image, target, label = self.grids[index] |
| if tile == -1: |
| return (image, *supervision(self.grids.classes, self.cells, target, any_label=label)) |
| uncertain = [] if self.keep_uncertain else record.get("uncertain", []) |
| if self.flip and random.random() < 0.5: |
| |
| image = image.flip(-1) |
| if tile is None: |
| n = int(label.numel() ** 0.5) |
| label = label.view(n, n).flip(-1).flatten() |
| uncertain = [c - c % n + n - 1 - c % n for c in uncertain] |
| labels, weights, any_labels, any_weights = supervision( |
| self.grids.classes, self.cells, target, cell_labels=label |
| ) |
| if tile is not None: |
| if tile in uncertain: |
| weights.zero_() |
| else: |
| for cell in uncertain: |
| weights[target, cell] = 0 |
| return image, labels, weights, any_labels, any_weights |
|
|
|
|
| class FolderDataset(Dataset): |
| """Single-label folder images: present for named targets, others assumed absent at low weight. |
| |
| A folder label says the image contains that target. Another target may also be visible, |
| so absence is only weakly supervised, except for background images, which contain none. |
| """ |
|
|
| def __init__(self, root, records, classes, size, cells=1, training=False, negative_weight=0.2): |
| self.root = Path(root) |
| self.records = records |
| self.classes = classes |
| self.cells = cells |
| self.negative_weight = negative_weight |
| self.transform = transform(size, training, noise=True) |
|
|
| def __len__(self): |
| return len(self.records) |
|
|
| def __getitem__(self, index): |
| record = self.records[index] |
| image = self.transform(read_rgb(self.root / record["path"])) |
| background = BACKGROUND in record["targets"] |
| any_weights = torch.full((len(self.classes),), 1.0 if background else self.negative_weight) |
| any_labels = torch.zeros(len(self.classes)) |
| for target in record["targets"]: |
| if target in self.classes: |
| any_labels[self.classes.index(target)] = any_weights[self.classes.index(target)] = 1 |
| empty = torch.zeros(len(self.classes), self.cells) |
| return image, empty, empty.clone(), any_labels, any_weights |
|
|
