Timechils/sapiens-mvp

1import os2import cv23import glob4import random5from tqdm import tqdm6 7# --- CONFIGURATION ---8# These paths point to your RAW YOLO data (Input)9IMAGES_DIR = "dataset_v2/ImageSet/train/images"10LABELS_DIR = "dataset_v2/ImageSet/train/labels"11 12# Output directory (The 9th Class)13OUTPUT_DIR = "dataset_v2_ready/No_Anomaly"14os.makedirs(OUTPUT_DIR, exist_ok=True)15 16TARGET_COUNT = 2500  # How many healthy images to create17CROP_SIZE = 224      # Match the ResNet-50 input size18 19def parse_yolo(label_path, img_w, img_h):20    """Reads YOLO labels and converts them to pixel coordinates."""21    boxes = []22    if not os.path.exists(label_path): return []23    with open(label_path, 'r') as f:24        for line in f:25            parts = line.strip().split()26            # YOLO format: class x_center y_center w h (normalized)27            x_c, y_c, w, h = map(float, parts[1:5])28            29            # Convert to pixel coordinates30            x1 = int((x_c - w/2) * img_w)31            y1 = int((y_c - h/2) * img_h)32            x2 = int((x_c + w/2) * img_w)33            y2 = int((y_c + h/2) * img_h)34            boxes.append((x1, y1, x2, y2))35    return boxes36 37def has_overlap(crop_box, fault_boxes):38    """Checks if the candidate crop touches any known fault."""39    cx1, cy1, cx2, cy2 = crop_box40    for (fx1, fy1, fx2, fy2) in fault_boxes:41        # Check for intersection42        if not (cx2 < fx1 or cx1 > fx2 or cy2 < fy1 or cy1 > fy2):43            return True # Overlap detected (Unsafe!)44    return False45 46def main():47    # Get list of all source images48    image_paths = glob.glob(os.path.join(IMAGES_DIR, "*.jpg"))49    print(f"🔍 Scanning {len(image_paths)} source images for safe zones...")50    51    count = 052    pbar = tqdm(total=TARGET_COUNT, desc="Generating Normal Crops")53 54    while count < TARGET_COUNT:55        # 1. Pick a random image56        img_path = random.choice(image_paths)57        label_path = os.path.join(LABELS_DIR, os.path.basename(img_path).replace(".jpg", ".txt"))58        59        img = cv2.imread(img_path)60        if img is None: continue61        62        h, w, _ = img.shape63        # Skip small images64        if h < CROP_SIZE or w < CROP_SIZE: continue65        66        # 2. Map out the 'Danger Zones' (Faults)67        faults = parse_yolo(label_path, w, h)68        69        # 3. Try to find a 'Safe Zone' (Background)70        # We try 20 times per image to find a spot with no faults71        for _ in range(20):72            # --- THE FIXED LINES ---73            rx = random.randint(0, w - CROP_SIZE)74            ry = random.randint(0, h - CROP_SIZE)75            # -----------------------76            77            crop_coords = (rx, ry, rx + CROP_SIZE, ry + CROP_SIZE)78            79            # If SAFE (no overlap with faults)80            if not has_overlap(crop_coords, faults):81                crop = img[ry:ry+CROP_SIZE, rx:rx+CROP_SIZE]82                83                # Check variance (ensure it's not just a pitch black corner)84                if crop.var() > 10:85                    save_name = f"normal_{count}_{os.path.basename(img_path)}"86                    cv2.imwrite(os.path.join(OUTPUT_DIR, save_name), crop)87                    88                    count += 189                    pbar.update(1)90                    break 91        92        if count >= TARGET_COUNT: break93 94    pbar.close()95    print(f"\n✅ Success! Generated {count} 'No_Anomaly' images in {OUTPUT_DIR}")96 97if __name__ == "__main__":98    main()99