helpers.py

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