File size: 6,646 Bytes
c8c00f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | #!/usr/bin/env python3
# scripts/ingest_public_data.py
# ============================================
# RUN ONCE: python scripts/ingest_public_data.py
# Populates ChromaDB with defect patch triplets
# ============================================
import chromadb
from sentence_transformers import SentenceTransformer
from PIL import Image
import json
from pathlib import Path
import argparse
# ============================================
# DATASET REGISTRY — Add new sources here
# ============================================
DATASETS = {
# ── Commercial-friendly datasets ──
"roboflow_manufacturing": {
"path": "data/external/roboflow_manufacturing",
"domain": "pcb",
"license": "cc-by-4.0",
"commercial_ok": True,
"attribution": "Roboflow Universe — Manufacturing Defect Detection",
"description": "PCB manufacturing defects, CC BY 4.0 licensed"
},
# ── Your proprietary data ──
"my_vcsel_captures": {
"path": "data/my_captures/vcsel",
"domain": "glass/vcsel",
"license": "proprietary",
"commercial_ok": True,
"attribution": "Internal",
"description": "VCSEL laser diode defect captures"
},
# ── Partner data (NDA) ──
# "partner_pcb_nda": {
# "path": "data/partners/pcb_client_a",
# "domain": "pcb",
# "license": "partner_nda",
# "commercial_ok": True,
# "attribution": "Partner NDA",
# "description": "PCB client defect images under NDA"
# },
}
def validate_triplet(folder: Path) -> dict | None:
"""Check if folder contains required patch triplet files."""
required = ["original_masked.png", "original_target.png", "artifact_target.png"]
paths = {f: folder / f for f in required}
for f, p in paths.items():
if not p.exists():
print(f" ⚠️ Missing {f} in {folder}, skipping")
return None
return {k: str(v) for k, v in paths.items()}
def generate_caption(folder_name: str, domain: str, dataset_desc: str) -> str:
"""Generate a text caption for embedding."""
# You can replace this with a VLM call for better captions
defect_type = folder_name.replace("_", " ")
return f"{defect_type} defect on {domain}. {dataset_desc}"
def ingest_dataset(collection, encoder, name: str, cfg: dict):
"""Ingest one dataset into ChromaDB."""
print(f"\n{'='*50}")
print(f"Dataset: {name}")
print(f"License: {cfg['license']} | Commercial: {'✅' if cfg['commercial_ok'] else '❌'}")
print(f"Path: {cfg['path']}")
print(f"{'='*50}")
base_path = Path(cfg["path"])
if not base_path.exists():
print(f" ⚠️ Path not found: {base_path}")
print(f" Create it and add patch triplet folders:")
print(f" {base_path}/defect_name_001/original_masked.png")
print(f" {base_path}/defect_name_001/original_target.png")
print(f" {base_path}/defect_name_001/artifact_target.png")
return 0
# NEW: Get set of IDs already in the collection to avoid re-processing
existing_ids = set(collection.get()["ids"])
count = 0
for triplet_folder in sorted(base_path.iterdir()):
if not triplet_folder.is_dir():
continue
# NEW: Skip if already processed
doc_id = f"{name}_{triplet_folder.name}"
if doc_id in existing_ids:
print(f" ⏭️ Skipping {triplet_folder.name} (already in DB)")
continue
paths = validate_triplet(triplet_folder)
if paths is None:
continue
caption = generate_caption(
triplet_folder.name,
cfg["domain"],
cfg.get("description", "")
)
embedding = encoder.encode(caption)
collection.upsert(
documents=[caption],
embeddings=[embedding.tolist()],
metadatas=[{
"paths": json.dumps(paths),
"domain": cfg["domain"],
"license": cfg["license"],
"commercial_ok": cfg["commercial_ok"],
"source": name,
"attribution": cfg["attribution"],
"defect_name": triplet_folder.name
}],
ids=[f"{name}_{triplet_folder.name}"]
)
count += 1
print(f" ✅ {triplet_folder.name}: {caption[:60]}...")
return count
def main():
parser = argparse.ArgumentParser(description="Ingest defect patch triplets into RAG DB")
parser.add_argument("--db-path", default="data/defect_db", help="ChromaDB persistent path")
parser.add_argument("--collection", default="defect_patches", help="Collection name")
parser.add_argument("--model", default="all-MiniLM-L6-v2", help="SentenceTransformer model")
args = parser.parse_args()
# Initialize DB
Path(args.db_path).mkdir(parents=True, exist_ok=True)
client = chromadb.PersistentClient(path=args.db_path)
# Delete existing collection if you want fresh start
# client.delete_collection(args.collection)
collection = client.get_or_create_collection(
name=args.collection,
metadata={"hnsw:space": "cosine"}
)
print(f"DB path: {args.db_path}")
print(f"Collection: {args.collection}")
print(f"Existing entries: {collection.count()}")
# Initialize encoder
print(f"\nLoading encoder: {args.model}")
encoder = SentenceTransformer(args.model)
# Ingest all datasets
total = 0
for name, cfg in DATASETS.items():
# Skip non-commercial datasets in commercial builds
if not cfg.get("commercial_ok", False):
print(f"\n⏭️ Skipping {name} — not commercial-friendly")
continue
count = ingest_dataset(collection, encoder, name, cfg)
total += count
print(f"\n{'='*50}")
print(f"TOTAL INGESTED: {total} patch triplets")
print(f"TOTAL IN DB: {collection.count()}")
print(f"{'='*50}")
# Print commercial summary
print("\n📋 Commercial License Summary:")
results = collection.get()
licenses = {}
for meta in results["metadatas"]:
lic = meta["license"]
licenses[lic] = licenses.get(lic, 0) + 1
for lic, count in licenses.items():
icon = "✅" if any(
cfg["license"] == lic and cfg.get("commercial_ok")
for cfg in DATASETS.values()
) else "❌"
print(f" {icon} {lic}: {count} entries")
if __name__ == "__main__":
main() |