import yaml import base64 import random import string import math import shutil import io import os import json from PIL import Image, ImageDraw, ImageFont from sahi.slicing import slice_image, slice_coco from sahi.utils.coco import Coco, CocoAnnotation, CocoImage, CocoCategory from pathlib import Path from ultralytics.data.utils import compress_one_image from ultralytics.utils.downloads import zip_directory from processor import * def generate_yolo_dataset_coco(num_images, output_directory="Data/yolo_dataset", size=(640, 640), allow_intersections=False, allow_grid=True, split_ratio=0.8, use_sahi_slicing=False, slice_size=640, overlap_ratio=0.2): if os.path.exists(output_directory): shutil.rmtree(output_directory) for split in ['train', 'val']: print(str(os.path.join(output_directory, split, "temp/images"))) # os.makedirs(os.path.join(output_directory, split, "images"), exist_ok=True) # os.makedirs(os.path.join(output_directory, split, "labels"), exist_ok=True) os.makedirs(os.path.join(output_directory, split, "temp/images"), exist_ok=True) os.makedirs(os.path.join(output_directory, split, "temp/outputs"), exist_ok=True) width, height = size # REDUCED SIZE: Changed multiplier from 0.15 to 0.07 (approx half size) shape_radius = (min(width, height) // 2) * 0.07 for n in range(num_images): img = Image.new("RGB", size, (255, 255, 255)) draw = ImageDraw.Draw(img) exclusion_zones = [] elements = [] current_split = 'train' if random.random() < split_ratio else 'val' # INCREASED QUANTITY: Targets, Grids, and Floating text counts increased to_place = [] num_targets = random.randint(3, 6) # More targets per image num_grids = random.randint(3, 5) if allow_grid else 0 num_empty_shapes = random.randint(4, 8) num_floating_texts = random.randint(5, 10) for _ in range(num_targets): to_place.append({"type": "target", "is_shape": True}) for _ in range(num_grids): to_place.append({"type": "grid", "is_shape": True}) for _ in range(num_empty_shapes): to_place.append({"type": "empty", "is_shape": True}) for _ in range(num_floating_texts): to_place.append({"type": "text", "is_shape": False}) for item in to_place: placed = False attempts = 0 while not placed and attempts < 300: # Increased attempts for denser packing tx, ty = random.randint(50, width-50), random.randint(50, height-50) # REDUCED BUFFER: Smaller exclusion radius to allow closer packing e_r = (shape_radius + 15) if item["is_shape"] else 15 if not any(math.sqrt((tx-z[0][0])**2 + (ty-z[0][1])**2) < (e_r + z[1]) for z in exclusion_zones): exclusion_zones.append(((tx, ty), e_r)) item["pos"] = (tx, ty) elements.append(item) placed = True attempts += 1 temp_img_name = f"img_{n}.png" temp_img_path = os.path.join(output_directory, current_split, "temp", "images", temp_img_name) coco = Coco() target_category = CocoCategory(id=0, name="target") coco.add_category(target_category) coco_image = CocoImage(file_name=temp_img_name, height=height, width=width) dynamic_font_size = int(shape_radius * 1.2) # Keeping font scaled to shape size try: font = ImageFont.load_default(size=dynamic_font_size) except: font = ImageFont.load_default() yolo_labels = [] # List to hold multiple target labels for el in elements: ex, ey = el["pos"] if el["type"] == "grid": grid_cell_size = shape_radius * 0.7 for row in range(2): for col in range(2): x1, y1 = ex + (col - 1) * grid_cell_size, ey + (row - 1) * grid_cell_size x2, y2 = x1 + grid_cell_size, y1 + grid_cell_size draw.rectangle([x1, y1, x2, y2], outline=(0, 0, 0), width=1) char = random.choice(string.ascii_uppercase + string.digits) draw.text(((x1 + x2)/2, (y1 + y2)/2), char, fill=(0, 0, 0), font=font, anchor="mm") elif el["is_shape"]: shape_type = random.choice(["diamond", "hexagon"]) if shape_type == "diamond": pts = [(ex, ey-shape_radius), (ex+shape_radius, ey), (ex, ey+shape_radius), (ex-shape_radius, ey)] else: pts = [(ex + shape_radius * math.cos(math.radians(i*60-30)), ey + shape_radius * math.sin(math.radians(i*60-30))) for i in range(6)] draw.polygon(pts, fill=(255, 255, 255), outline=(0, 0, 0), width=2) if el["type"] == "target": txt = f"{random.choice(string.ascii_uppercase)}{random.randint(0, 9)}" draw.text((ex, ey), txt, fill=(0, 0, 0), font=font, anchor="mm") all_x, all_y = [p[0] for p in pts], [p[1] for p in pts] xmin, ymin, xmax, ymax = min(all_x), min(all_y), max(all_x), max(all_y) x_center, y_center = ((xmin + xmax) / 2) / width, ((ymin + ymax) / 2) / height w_norm, h_norm = (xmax - xmin) / width, (ymax - ymin) / height yolo_labels.append(f"0 {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}") # draw.rectangle([xmin, ymin, xmax, ymax], outline=(255, 0, 0), width=2) bbox = [ex - shape_radius, ey - shape_radius, shape_radius * 2, shape_radius * 2] coco_image.add_annotation( CocoAnnotation( bbox=bbox, category_id=0, category_name="target" ) ) else: txt = f"{random.choice(string.ascii_uppercase)}{random.randint(0, 9)}" draw.text((ex, ey), txt, fill=(0, 0, 0), font=font, anchor="mm") # INCREASED BACKGROUND NOISE: Changed from 12 to 30 lines lines_drawn = 0 while lines_drawn < 30: p1, p2 = (random.randint(0, width), random.randint(0, height)), (random.randint(0, width), random.randint(0, height)) random_width = random.randint(1, 3) # Slightly thinner noise lines can_draw = True if not allow_intersections: for z_pos, z_r in exclusion_zones: dx, dy = p2[0]-p1[0], p2[1]-p1[1] mag_sq = dx**2 + dy**2 + 1e-9 u = max(0, min(1, ((z_pos[0]-p1[0])*dx + (z_pos[1]-p1[1])*dy) / mag_sq)) if ((p1[0] + u*dx - z_pos[0])**2 + (p1[1] + u*dy - z_pos[1])**2) < z_r**2: can_draw = False; break if can_draw: draw.line([p1, p2], fill=(0, 0, 0), width=random_width) lines_drawn += 1 # --- SAHI SLICING LOGIC --- img.save(temp_img_path) coco.add_image(coco_image) source_json_path = os.path.join(output_directory, current_split, "source_coco.json") with open(source_json_path, 'w') as f: json.dump(coco.json, f) if use_sahi_slicing: # slice_image calculates coordinates for each tile _, sliced_path_result = slice_coco( coco_annotation_file_path=source_json_path, # Pass the string path image_dir=os.path.join(output_directory, current_split, "temp", "images"), output_dir=os.path.join(output_directory, current_split, "temp", "images"), output_coco_annotation_file_name="temp_sliced_meta.json", slice_height=slice_size, slice_width=slice_size, overlap_height_ratio=overlap_ratio, overlap_width_ratio=overlap_ratio, min_area_ratio=0.1, verbose=False ) # Export the results specifically to YOLO .txt format # slice_results['coco'] contains the newly calculated slices final_path = str(sliced_path_result) final_coco = Coco.from_coco_dict_or_path( final_path, image_dir=os.path.join(output_directory, current_split, "temp", "images") ) # 4. Export to YOLO # SAHI will automatically create 'images' and 'labels' folders inside this output_dir final_coco.export_as_yolo( output_dir=os.path.join(output_directory, current_split, "temp", "outputs"), disable_symlink=True ) else: # Manual Save if slicing is off img.save(os.path.join(output_directory, current_split, "images", f"img_{n:04d}.png")) # (Label logic for non-sliced omitted for brevity, similar to previous steps) if os.path.exists(temp_img_path): os.remove(temp_img_path) if use_sahi_slicing: for current_split in ['train', 'val']: print(current_split) sahi_export_dir = os.path.join(output_directory, current_split, "temp") # --- REORGANIZATION SECTION --- # Define where SAHI put things vs where we want them sahi_outputs = os.path.join(sahi_export_dir, "outputs", "train") os.makedirs(os.path.join(output_directory, current_split, "images"), exist_ok=True) os.makedirs(os.path.join(output_directory, current_split, "labels"), exist_ok=True) target_images_dir = os.path.join(output_directory, current_split, "images") target_labels_dir = os.path.join(output_directory, current_split, "labels") # 1. Move sliced images from .../labels/images to .../images if os.path.exists(sahi_outputs): for f in os.listdir(sahi_outputs): if ".png" in f: shutil.move(os.path.join(sahi_outputs, f), os.path.join(target_images_dir, f)) # 2. Move .txt files from .../labels/labels to .../labels if os.path.exists(sahi_outputs): for f in os.listdir(sahi_outputs): if ".txt" in f: shutil.move(os.path.join(sahi_outputs, f), os.path.join(target_labels_dir, f)) if os.path.exists(os.path.join(sahi_export_dir, "outputs")): for f in os.listdir(os.path.join(sahi_export_dir, "outputs")): if not os.path.exists(os.path.join(output_directory, "data.yml")): if ".yml" in f: shutil.move(os.path.join(os.path.join(sahi_export_dir, "outputs"), f), output_directory) # 2. Edit the contents # Assuming your YAML has a 'training' key or you want to add a new one data = {} data['path'] = os.path.abspath(output_directory) data['train'] = "train/images" data['val'] = "val/images" data['names'] = {0: "target_shape"} # 3. Save as a .yaml file with open(os.path.join(output_directory, 'data.yaml'), 'w') as file: # default_flow_style=False keeps the readable block format yaml.dump(data, file, default_flow_style=False, sort_keys=False) # 3. Delete the now-empty nested 'labels' folder structure created by SAHI shutil.rmtree(sahi_export_dir) # --- CLEANUP --- # Remove original high-res image if os.path.exists(temp_img_path): os.remove(temp_img_path) # Remove the source JSON we created if os.path.exists(source_json_path): os.remove(source_json_path) # Remove the intermediate sliced JSON if os.path.exists(sliced_path_result): os.remove(sliced_path_result) if os.path.exists(sahi_export_dir): os.remove(sahi_export_dir) if os.path.exists(os.path.join(output_directory, current_split, "source_coco.json")): os.remove(os.path.join(output_directory, current_split, "source_coco.json")) if os.path.exists(os.path.join(output_directory, "data.yml")): os.remove(os.path.join(output_directory, "data.yml")) print(f"Dataset generated at {output_directory}") # Generate yolo specific datasets def generate_yolo_dataset_original(num_images, output_dir="Data/yolo_dataset", size=(640, 640), allow_intersections=False, allow_grid=True, split_ratio=0.8): # 1. Setup Directory Structure for YOLO splits for split in ['train', 'val']: os.makedirs(os.path.join(output_dir, split, "images"), exist_ok=True) os.makedirs(os.path.join(output_dir, split, "labels"), exist_ok=True) width, height = size shape_radius = (min(width, height) // 2) * 0.15 for n in range(num_images): # --- [Generation Logic Remains the Same] --- img = Image.new("RGB", size, (255, 255, 255)) draw = ImageDraw.Draw(img) exclusion_zones = [] elements = [] # Decide if this specific image goes to 'train' or 'val' current_split = 'train' if random.random() < split_ratio else 'val' # Determine items to place num_floating_texts = random.randint(1, 2) if allow_grid: to_place = [{"type": "target", "is_shape": True}, {"type": "grid", "is_shape": True}] else: to_place = [{"type": "target", "is_shape": True}] for _ in range(2): to_place.append({"type": "empty", "is_shape": True}) for _ in range(num_floating_texts): to_place.append({"type": "text", "is_shape": False}) # Placement Loop for item in to_place: placed = False attempts = 0 while not placed and attempts < 200: tx, ty = random.randint(100, width-100), random.randint(100, height-100) e_r = (shape_radius + 50) if item["is_shape"] else 50 if not any(math.sqrt((tx-z[0][0])**2 + (ty-z[0][1])**2) < (e_r + z[1]) for z in exclusion_zones): exclusion_zones.append(((tx, ty), e_r)) item["pos"] = (tx, ty) elements.append(item) placed = True attempts += 1 # Drawing Logic (Simplified for brevity, use previous detailed drawing code here) yolo_label_data = "" dynamic_font_size = int(shape_radius) try: font = ImageFont.load_default(size=dynamic_font_size) except: font = ImageFont.load_default() for el in elements: ex, ey = el["pos"] # --- NEW GRID LOGIC --- if el["type"] == "grid": grid_cell_size = shape_radius * 0.8 for row in range(2): for col in range(2): # Calculate cell corners x1 = ex + (col - 1) * grid_cell_size y1 = ey + (row - 1) * grid_cell_size x2 = x1 + grid_cell_size y2 = y1 + grid_cell_size # Draw cell draw.rectangle([x1, y1, x2, y2], outline=(0, 0, 0), width=2) # Add random char to cell char = random.choice(string.ascii_uppercase + string.digits) draw.text(((x1 + x2)/2, (y1 + y2)/2), char, fill=(0, 0, 0), font=font, anchor="mm") # --- EXISTING SHAPE LOGIC --- elif el["is_shape"]: shape_type = random.choice(["diamond", "hexagon"]) if shape_type == "diamond": pts = [(ex, ey-shape_radius), (ex+shape_radius, ey), (ex, ey+shape_radius), (ex-shape_radius, ey)] else: pts = [(ex + shape_radius * math.cos(math.radians(i*60-30)), ey + shape_radius * math.sin(math.radians(i*60-30))) for i in range(6)] draw.polygon(pts, fill=(255, 255, 255), outline=(0, 0, 0), width=3) if el["type"] == "target": txt = f"{random.choice(string.ascii_uppercase)}{random.randint(0, 9)}" draw.text((ex, ey), txt, fill=(0, 0, 0), font=font, anchor="mm") all_x, all_y = [p[0] for p in pts], [p[1] for p in pts] xmin, ymin, xmax, ymax = min(all_x)-5, min(all_y)-5, max(all_x)+5, max(all_y)+5 x_center, y_center = ((xmin + xmax) / 2) / width, ((ymin + ymax) / 2) / height w_norm, h_norm = (xmax - xmin) / width, (ymax - ymin) / height yolo_label_data = f"0 {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}" draw.rectangle([xmin, ymin, xmax, ymax], outline=(255, 0, 0), width=2) else: # Floating text txt = f"{random.choice(string.ascii_uppercase)}{random.randint(0, 9)}" draw.text((ex, ey), txt, fill=(0, 0, 0), font=font, anchor="mm") # Background lines (Checks against exclusion_zones, which now includes the grid) lines_drawn = 0 while lines_drawn < 12: p1, p2 = (random.randint(0, width), random.randint(0, height)), (random.randint(0, width), random.randint(0, height)) random_width = random.randint(1, 5) can_draw = True if not allow_intersections: for z_pos, z_r in exclusion_zones: dx, dy = p2[0]-p1[0], p2[1]-p1[1] mag_sq = dx**2 + dy**2 + 1e-9 u = max(0, min(1, ((z_pos[0]-p1[0])*dx + (z_pos[1]-p1[1])*dy) / mag_sq)) if ((p1[0] + u*dx - z_pos[0])**2 + (p1[1] + u*dy - z_pos[1])**2) < z_r**2: can_draw = False; break if can_draw: draw.line([p1, p2], fill=(0, 0, 0), width=random_width) lines_drawn += 1 # --- Save to the determined split folder --- file_base = f"img_{n:04d}" # Save Image img_path = os.path.join(output_dir, current_split, "images", f"{file_base}.png") img.save(img_path) # Save Label (only if target was generated/found) if yolo_label_data: lbl_path = os.path.join(output_dir, current_split, "labels", f"{file_base}.txt") with open(lbl_path, "w") as f: f.write(yolo_label_data) # 4. Create data.yaml yaml_content = f"path: {os.path.abspath(output_dir)}\ntrain: images\nval: images\n\nnames:\n 0: target_shape" with open(os.path.join(output_dir, "data.yaml"), "w") as f: f.write(yaml_content) print(f"Dataset created at {output_dir}. Ready for YOLO training.")