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(num_images, output_directory="Data/yolo_dataset", size=(640, 640), allow_intersections=False, allow_grid=True, split_ratio=0.8): if os.path.exists(output_directory): shutil.rmtree(output_directory) for split in ['train', 'val']: os.makedirs(os.path.join(output_directory, split, "images"), exist_ok=True) os.makedirs(os.path.join(output_directory, split, "labels"), 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): print(f"Generating image {n+1}/{num_images}", end='\r') 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 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 = [] 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)-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_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) 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 file_base = f"img_{n:04d}" # Save Image img_path = os.path.join(output_directory, current_split, "images", f"{file_base}.png") img.save(img_path) # Save Label (only if target was generated/found) if len(yolo_labels): lbl_path = os.path.join(output_directory, current_split, "labels", f"{file_base}.txt") with open(lbl_path, "w") as f: f.write('\n'.join(yolo_labels)) # 4. Create data.yaml yaml_content = f"path: {os.path.abspath(output_directory)}\ntrain: train/images\nval: val/images\n\nnames:\n 0: target_shape" with open(os.path.join(output_directory, "data.yaml"), "w") as f: f.write(yaml_content) print(f"Dataset generated at {output_directory}") def binary_invert_and_save(image_path, threshold=200): # 1. Load the image img = cv2.imread(image_path) if img is None: print(f"Error: Could not read image at {image_path}") return # 2. Convert to grayscale gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 3. Apply Threshold (White background, Black objects) # Pixels > threshold become 255 (White), others become 0 (Black) _, binary = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY) # 4. Extract original filename and extension original_name = os.path.basename(image_path) # Optional: Save to a 'processed' folder to keep things organized output_dir = "processed_images" if not os.path.exists(output_dir): os.makedirs(output_dir) save_path = os.path.join(output_dir, original_name) # 5. Save the file cv2.imwrite(save_path, binary) print(f"Saved: {save_path}") def test_cv_model(file, model, output_dir): images = convert_from_bytes(open(file, "rb").read()) for i, image in enumerate(images): # image_bytes = image["content"] # original_image = Image.open(io.BytesIO(image)) # 2. Run Inference # model(original_image) typically returns a list of Results prediction_results = model(image) # 3. Extract the first result (since we passed one image) # The .plot() method draws ALL detected boxes onto the frame automatically result = prediction_results[0] annotated_frame = result.plot() # 4. Convert BGR numpy array (standard for CV2) to RGB for PIL # [..., ::-1] reverses the last dimension from BGR to RGB annotated_img_pil = Image.fromarray(annotated_frame[..., ::-1]) annotated_img_pil.save(os.path.join(output_dir, f"annotated_page_{i+1}.jpg")) # 5. Return the final edited PIL Image object return annotated_img_pil