archive.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_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.")