# === REVEL → KITTI dataset (train/val), safe jitter (objects fixed) ===
import omni.replicator.core as rep, os
from pathlib import Path

USER     = os.getenv("USER")
OUT_ROOT = Path(f"/home/{USER}/datasets/revel_kitti")  # Output directory
RES      = (1280, 720)                                 # Image resolution — HD standard
N_TOTAL  = 6000                                        # Total frame count
SPLIT    = 0.9                                         # Train ratio (90%)
N_TRAIN  = int(N_TOTAL * SPLIT)
N_VAL    = N_TOTAL - N_TRAIN

# 1) Target: All objects with Semantics.class tag
targets = rep.get.prims(semantics=[('class','*')])
assert targets, "ERROR: No object found with 'Semantics.class' tag. Add it to Root Xform!"

# 2) Camera + Render Product — FOV and gaze point approximating human eye level
cam = rep.create.camera(position=(0,1.6,2.2), look_at=(0,0.6,0), focal_length=24.0)
rp  = rep.create.render_product(cam, RES)

# 3) JITTER ONLY LIGHT AND CAMERA — DO NOT MOVE OBJECTS!
def jitter():
    # Dome light: ambient illumination — randomized temperature and intensity
    rep.randomizer.light(light_type='dome',
                         intensity=rep.distribution.uniform(800, 3000),
                         temperature=rep.distribution.uniform(4500, 8500))
    # Sphere light: directional spot effect — randomized position and count
    rep.create.light(light_type='sphere',
                     intensity=rep.distribution.uniform(1000, 4000),
                     position =rep.distribution.uniform((-2,2,-2),(2,4,2)),
                     temperature=rep.distribution.uniform(3000, 9000),
                     count    =rep.distribution.choice([1,2,3]))
    # Slightly perturb camera — focusing on objects
    with cam:
        rep.modify.pose(
            position=rep.distribution.uniform((-1.5,1.2,-1.5),(1.5,2.2,1.5)),  # +/- 1.5m lateral
            look_at =rep.distribution.choice(targets)                          # Look at random object
        )

rep.randomizer.register(jitter)  # Register randomizer

writer = rep.WriterRegistry.get("KittiWriter")  # Write in standard KITTI format

def run_split(out_dir, frames):
    writer.initialize(output_dir=str(out_dir))   # Set output directory
    writer.attach([rp])                          # Attach render product
    print(f"[RUN] {out_dir}  frames={frames}")
    rep.orchestrator.run(num_frames=frames)      # Generate frames

# Generate train and val sets separately
run_split(OUT_ROOT / "train", N_TRAIN)
run_split(OUT_ROOT / "val",   N_VAL)
print(f"[DONE] KITTI dataset ready → {OUT_ROOT}")

import os, cv2, shutil

KITTI_ROOT = os.path.expanduser("~/datasets/revel_kitti")
YOLO_ROOT  = os.path.expanduser("~/datasets/revel_yolo")
CLASSES_TXT= os.path.expanduser("~/classes.txt")  # One class name per line
IMG_EXT    = (".png", ".jpg", ".jpeg")

# Load class names and map to IDs
classes = [l.strip() for l in open(CLASSES_TXT) if l.strip()]
name2id = {n:i for i,n in enumerate(classes)}  # "bottle" → 0, "can" → 1, ...

def find_label_for(stem, search_dir):
    """Find KITTI label file — usually in same directory"""
    for root, _, files in os.walk(search_dir):
        if stem + ".txt" in files:
            return os.path.join(root, stem + ".txt")
    return None

def convert_split(split):
    split_dir = os.path.join(KITTI_ROOT, split)
    out_img = os.path.join(YOLO_ROOT, "images", split)
    out_lbl = os.path.join(YOLO_ROOT, "labels", split)
    os.makedirs(out_img, exist_ok=True)
    os.makedirs(out_lbl, exist_ok=True)

    for root, _, files in os.walk(split_dir):
        for fn in files:
            if not fn.lower().endswith(IMG_EXT):
                continue
            img_path = os.path.join(root, fn)
            stem = os.path.splitext(fn)[0]
            img = cv2.imread(img_path)
            if img is None:
                continue
            h, w = img.shape[:2]  # Image dimensions — required for normalization

            lab_path = find_label_for(stem, split_dir)
            yolo_lines = []
            if lab_path and os.path.exists(lab_path):
                with open(lab_path) as f:
                    for line in f:
                        parts = line.strip().split()
                        if len(parts) < 8:
                            continue
                        cls_name = parts[0]  # Class name (e.g., "bottle")
                        try:
                            # KITTI bbox format: xmin, ymin, xmax, ymax (in pixels)
                            xmin, ymin, xmax, ymax = map(float, parts[4:8])
                        except:
                            # Fallback: extract numeric values
                            nums = [float(p) for p in parts if p.replace('.','',1).isdigit()]
                            if len(nums) >= 4:
                                xmin, ymin, xmax, ymax = nums[:4]
                            else:
                                continue
                        # Normalize: center x,y and width/height ratios
                        cx = (xmin + xmax) / 2.0 / w
                        cy = (ymin + ymax) / 2.0 / h
                        bw = (xmax - xmin) / w
                        bh = (ymax - ymin) / h
                        cid = name2id.get(cls_name)
                        # Add if valid class and valid bbox size
                        if cid is not None and 0 < bw <= 1 and 0 < bh <= 1:
                            yolo_lines.append(f"{cid} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")

            # Copy image and write label to target directory
            dst_img = os.path.join(out_img, fn)
            shutil.copy2(img_path, dst_img)
            with open(os.path.join(out_lbl, stem + ".txt"), "w") as f:
                f.write("\n".join(yolo_lines))

# Convert both train and val splits
for sp in ("train","val"):
    convert_split(sp)

print(f"✅ YOLO dataset ready → {YOLO_ROOT}")