23_inference_segmenter.py

import os
import json
from ultralytics import YOLO
from PIL import Image, ImageDraw
import numpy as np
import tifffile
import csv
import random
import colorsys
import argparse
import logging
import sys
import shutil
from tqdm import tqdm

# --- Default Configuration ---
DEFAULT_INPUT_DIR = "inference_in"
DEFAULT_OUTPUT_DIR = "inference_out"
DEFAULT_LOG_FILE = "master_log.txt"
DEFAULT_LEAF_INCLUSION_MODE = 'touching_or_inside'
SAMPLE_MODEL_BASE_DIR = os.path.join("models", "pretrained_model_data", "segmentation", "model")
SAMPLE_DATA_PATH = "models/pretrained_model_data/segmentation/example_inference_in"

def setup_logging(log_file):
    """Configures logging to output to both console and file."""
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    if logger.hasHandlers():
        logger.handlers.clear()

    # File handler for master_log.txt
    file_handler = logging.FileHandler(log_file, mode='a', encoding="utf-8")
    file_handler.setLevel(logging.INFO)
    file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)

    # Console handler for INFO and above
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_formatter = logging.Formatter('%(levelname)s: %(message)s')
    console_handler.setFormatter(console_formatter)
    logger.addHandler(console_handler)

class DynamicColorManager:
    """
    Manages color assignments for object detection classes.
    All instances of a class share the same hue, but each instance gets a different saturation.
    """
    def __init__(self):
        self._class_hues = {}
        self._class_instance_count = {}

    def _get_hue_for_class(self, class_name):
        """Assigns a fixed hue per class, deterministically based on class name."""
        if class_name not in self._class_hues:
            # Use a hash of the class name to get a deterministic hue between 0 and 1
            h = abs(hash(class_name)) % 360
            self._class_hues[class_name] = h / 360.0
            self._class_instance_count[class_name] = 0
        return self._class_hues[class_name]

    def get_color(self, class_name):
        """
        Gets a color for a given class name and instance.
        All instances of a class share the same hue, but each gets a different saturation.
        """
        hue = self._get_hue_for_class(class_name)
        idx = self._class_instance_count[class_name]
        saturations = [0.4, 0.55, 0.7, 0.85, 1.0, 0.25]
        saturation = saturations[idx % len(saturations)]
        lightness = 0.55
        self._class_instance_count[class_name] += 1
        rgb_float = colorsys.hls_to_rgb(hue, lightness, saturation)
        return tuple(int(c * 255) for c in rgb_float)

def load_plot_border_mask(image_filename, plot_border_dir, image_width, image_height):
    """Loads the plot border polygon from a LabelMe JSON file and returns a binary mask and the polygon points."""
    base_filename, _ = os.path.splitext(image_filename)
    json_path = os.path.join(plot_border_dir, f"{base_filename}.json")
    
    points = None

    if not os.path.exists(json_path):
        logging.info(f"No plot border file found at {json_path}. Assuming entire image is valid.")
        return Image.new('L', (image_width, image_height), 255), None

    try:
        with open(json_path, 'r') as f:
            data = json.load(f)
    except Exception as e:
        logging.error(f"Error reading or parsing plot border file {json_path}: {e}")
        return None, None

    plot_shape = next((s for s in data.get('shapes', []) if s.get('label', '').lower() == 'plot'), None)

    if not plot_shape:
        logging.warning(f"No shape with label 'plot' found in {json_path}. Assuming entire image is valid.")
        return Image.new('L', (image_width, image_height), 255), None

    mask = Image.new('L', (image_width, image_height), 0)
    draw = ImageDraw.Draw(mask)
    
    raw_points = plot_shape['points']
    if plot_shape.get('shape_type', 'polygon').lower() == 'rectangle' and len(raw_points) == 2:
        p1, p2 = raw_points[0], raw_points[1]
        min_x, max_x = min(p1[0], p2[0]), max(p1[0], p2[0])
        min_y, max_y = min(p1[1], p2[1]), max(p1[1], p2[1])
        points = [(min_x, min_y), (max_x, min_y), (max_x, max_y), (min_x, max_y)]
    else:
        points = [tuple(p) for p in raw_points]

    if len(points) > 1:
        draw.polygon(points, outline=255, fill=255)
    else:
        logging.warning(f"Not enough points to draw polygon from {json_path}.")
        return None, None

    return mask, points

def find_latest_model(search_dir="runs"):
    """Finds the latest 'best.pt' model file in a directory."""
    latest_model_path = None
    latest_model_time = 0
    if not os.path.isdir(search_dir):
        logging.error(f"Model search directory '{search_dir}' not found.")
        return None
    
    for root, _, files in os.walk(search_dir):
        if 'best.pt' in files:
            model_file_path = os.path.join(root, 'best.pt')
            try:
                model_time = os.path.getmtime(model_file_path)
                if model_time > latest_model_time:
                    latest_model_time = model_time
                    latest_model_path = model_file_path
            except OSError as e:
                logging.warning(f"Could not access {model_file_path}: {e}")
    return latest_model_path

def main():
    parser = argparse.ArgumentParser(
        description="Run YOLO segmentation inference on a directory of images.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--model-path", type=str, default=None,
                        help="Path to the YOLO model file (.pt). If not provided, finds the latest 'best.pt' in 'runs'.")
    parser.add_argument("--input-dir", type=str, default=DEFAULT_INPUT_DIR,
                        help="Directory containing input images for inference.")
    parser.add_argument("--output-dir", type=str, default=DEFAULT_OUTPUT_DIR,
                        help="Base directory to save all output files and subdirectories.")
    parser.add_argument("--leaf-inclusion-mode", type=str, default=DEFAULT_LEAF_INCLUSION_MODE,
                        choices=['fully_inside', 'touching_or_inside'],
                        help="Defines how to count leaves relative to the plot border.")
    parser.add_argument("--use-sample-model", action="store_true",
                        help="Use a pretrained model from the 'models/pretrained_model_data' directory. Use with --sample-model-size.")
    parser.add_argument("--use-sample-images", action="store_true",
                        help=f"Use the sample images and labels from '{SAMPLE_DATA_PATH}'. Overrides --input-dir.")
    parser.add_argument("--sample-model-size", type=str, default='x',
                        choices=['n', 's', 'm', 'l', 'x'],
                        help="YOLO model size to use with --use-sample-model. Options: 'n', 's', 'm', 'l', 'x'.")
    parser.add_argument("--log-file", type=str, default=DEFAULT_LOG_FILE,
                        help="Path to the master log file for appending logs.")
    parser.add_argument("--no-clean", action="store_true",
                        help="Do not clean the output directory before running.")

    args = parser.parse_args()

    setup_logging(args.log_file)
    logging.info("--- Starting script: 23_inference_segmenter.py ---")

    model_path = args.model_path
    input_dir = args.input_dir
    
    # --- Path Configuration based on arguments ---
    # 1. Handle Model Path
    if args.use_sample_model:
        logging.info(f"Using sample model mode with size '{args.sample_model_size}'.")
        model_path = os.path.join(SAMPLE_MODEL_BASE_DIR, f"yolo11{args.sample_model_size}_seg_eval", "weights", "best.pt")
        if args.model_path:
            logging.warning(f"Overriding explicit --model-path '{args.model_path}' with sample model.")
    
    # 2. Handle Input Data Path
    if args.use_sample_images:
        logging.info(f"Using sample images. Overriding input directory.")
        input_dir = SAMPLE_DATA_PATH
    
    # Determine the plot border directory based on whether sample images are used
    if args.use_sample_images:
        # For sample data, the plot borders are in a specific subdirectory
        plot_border_dir = os.path.join(input_dir, "plot_border")
    else:
        # For user-provided data, assume 'plot_border' is a direct subdirectory of the input
        plot_border_dir = os.path.join(input_dir, "plot_border")

    # Log the number of plot border annotations found
    if os.path.isdir(plot_border_dir):
        border_files = [f for f in os.listdir(plot_border_dir) if f.lower().endswith('.json')]
        logging.info(f"Found {len(border_files)} plot border annotations in '{plot_border_dir}'.")
        if not border_files:
            logging.warning("No plot border annotation files (.json) found. All images will be treated as fully valid.")
    else:
        logging.warning(f"Plot border directory not found at '{plot_border_dir}'. All images will be treated as fully valid.")

    # --- Final Path Validation ---
    if not model_path:
        logging.info("No model path specified or selected, searching for the latest model in 'runs' directory...\nYou can also use the '--use-sample-model' argument to use a pretrained sample model.")
        model_path = find_latest_model()
        if not model_path:
            logging.error("No 'best.pt' model found. Please specify a model with --model-path or --use-sample-model.")
            sys.exit(1)
    
    if not os.path.exists(model_path):
        logging.error(f"Model file not found at '{model_path}'.")
        sys.exit(1)
    
    if not os.path.isdir(input_dir):
        logging.error(f"Input directory not found at '{input_dir}'.")
        sys.exit(1)

    logging.info(f"Using model: {model_path}")
    logging.info(f"Input directory: {input_dir}")
    logging.info(f"Output directory: {args.output_dir}")

    # --- Define and prepare output directories ---
    cutouts_outline_dir = os.path.join(args.output_dir, "cutouts_with_outline")
    cutouts_leaf_only_dir = os.path.join(args.output_dir, "cutouts_leaf_only")
    cutouts_masked_dir = os.path.join(args.output_dir, "cutouts_masked")
    binary_masks_dir = os.path.join(args.output_dir, "binary_masks")
    num_result_dir = os.path.join(args.output_dir, "numeric")
    frame_dir = os.path.join(args.output_dir, "frames")

    for d in [cutouts_outline_dir, cutouts_leaf_only_dir, cutouts_masked_dir, binary_masks_dir, num_result_dir, frame_dir]:
        os.makedirs(d, exist_ok=True)

    if not args.no_clean:
        logging.info(f"Cleaning output directory: {args.output_dir}")
        items_to_delete = []
        dirs_to_clean = [cutouts_outline_dir, cutouts_leaf_only_dir, cutouts_masked_dir, binary_masks_dir, num_result_dir, frame_dir]
        for d in dirs_to_clean:
            if not os.path.isdir(d):
                continue
            for item in os.listdir(d):
                if item.lower() != ".gitkeep":
                    items_to_delete.append(os.path.join(d, item))
        
        if items_to_delete:
            for item_path in tqdm(items_to_delete, desc="Cleaning files", unit="file", ncols=100):
                try:
                    if os.path.isdir(item_path):
                        shutil.rmtree(item_path)
                    else:
                        os.remove(item_path)
                except Exception as e:
                    logging.error(f"Error removing {item_path}: {e}")

    # --- CSV Setup ---
    csv_filepath = os.path.join(num_result_dir, "results.csv")
    csv_header = [
        "image_name", "image_title", "original_width", "original_height", "instance_number",
        "bbox_center_x", "bbox_center_y", "bbox_width", "bbox_height",
        "confidence", "class_id", "class", "polygon_coordinates"
    ]

    csv_writer = None # Initialize csv_writer to None
    # Open CSV file for writing
    try:
        csv_file = open(csv_filepath, 'w', newline='')
        csv_writer = csv.writer(csv_file)
        csv_writer.writerow(csv_header)
        print(f"Opened CSV file for writing: {csv_filepath}")
    except IOError as e:
        logging.error(f"Error opening CSV file {csv_filepath}: {e}")
        return

    # --- Model Loading ---
    model = YOLO(model_path)
    logging.info("YOLO model loaded successfully.")

    # --- Main Processing Loop ---
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.tiff')]
    total_files = len(image_files)
    for i, filename in enumerate(image_files):
        input_path = os.path.join(input_dir, filename)
        logging.info(f"Processing image ({i + 1} of {total_files}): {filename}")
        
        image_title_str = "N/A"
        try:
            with tifffile.TiffFile(input_path) as tif:
                page = tif.pages[0]
                if 'ImageDescription' in page.tags:
                    value = page.tags['ImageDescription'].value
                    image_title_str = value.decode(errors='replace').strip() if isinstance(value, bytes) else str(value).strip()
        except Exception as e:
            logging.warning(f"Could not read image title metadata for {filename}: {e}")

        try:
            original_pil_image = Image.open(input_path).convert('RGB')
        except Exception as e:
            logging.error(f"Error opening image {input_path}: {e}")
            continue
            
        img_width, img_height = original_pil_image.size
        plot_border_mask_pil, plot_border_points = load_plot_border_mask(filename, plot_border_dir, img_width, img_height)
        if plot_border_mask_pil is None:
            logging.error(f"Skipping image {filename} due to error loading plot border.")
            continue
        
        plot_border_mask_np = (np.array(plot_border_mask_pil) > 0).astype(np.uint8)
        
        annotated_image_for_frame = original_pil_image.copy()
        draw_on_frame = ImageDraw.Draw(annotated_image_for_frame, 'RGBA')
        color_manager = DynamicColorManager()

        if plot_border_points:
            draw_on_frame.polygon(plot_border_points, outline="cyan", width=5)
        
        results = model(original_pil_image, verbose=False)
        total_detections = 0
        valid_instances = 0

        for result in results:
            if not (result.masks and result.boxes):
                continue

            total_detections += len(result.masks)
            for i in range(len(result.masks)):
                box, mask_item = result.boxes[i], result.masks[i]

                instance_mask_pil = Image.new('L', (img_width, img_height), 0)
                # Check for valid polygon coordinates
                if not (mask_item.xy and len(mask_item.xy) > 0 and len(mask_item.xy[0]) >= 2):
                    logging.warning(f"Instance {i} in image {filename} has an invalid or empty polygon. Skipping this instance.")
                    continue
                try:
                    ImageDraw.Draw(instance_mask_pil).polygon([tuple(p) for p in mask_item.xy[0]], fill=1)
                except Exception as e:
                    logging.warning(f"Error drawing polygon for instance {i} in image {filename}: {e}. Skipping this instance.")
                    continue
                
                instance_mask_np = np.array(instance_mask_pil)
                is_overlapping = (instance_mask_np == 1) & (plot_border_mask_np == 1)
                
                skip_instance = False
                if args.leaf_inclusion_mode == 'fully_inside' and np.any((instance_mask_np == 1) & (plot_border_mask_np == 0)):
                    skip_instance = True
                elif args.leaf_inclusion_mode == 'touching_or_inside' and not np.any(is_overlapping):
                    skip_instance = True
                
                if skip_instance:
                    continue
                
                valid_instances += 1
                class_id = int(box.cls[0])
                class_name = model.names.get(class_id, "unknown")
                instance_color = color_manager.get_color(class_name)
                
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                polygon_pixels = mask_item.xy[0].tolist() if mask_item.xy else []
                
                draw_on_frame.polygon(polygon_pixels, fill=instance_color + (100,), outline=instance_color + (255,))
                
                polygon_str = ";".join([f"{p[0]:.4f},{p[1]:.4f}" for p in polygon_pixels])
                csv_writer.writerow([
                    filename, image_title_str, img_width, img_height, i + 1,
                    (x1 + x2) / 2, (y1 + y2) / 2, x2 - x1, y2 - y1,
                    box.conf[0].item(), class_id, class_name, polygon_str
                ])
                
                padding = 5
                x1_p, y1_p, x2_p, y2_p = max(0, x1 - padding), max(0, y1 - padding), min(img_width, x2 + padding), min(img_height, y2 + padding)
                
                # Save cutouts
                base_fn, _ = os.path.splitext(filename)
                cutout_fn = f"{base_fn}_{i}.tiff"
                
                # Cutout with outline
                cutout_pil = original_pil_image.crop((x1_p, y1_p, x2_p, y2_p))
                draw_cutout = ImageDraw.Draw(cutout_pil)
                poly_cutout = [(px - x1_p, py - y1_p) for px, py in polygon_pixels]
                draw_cutout.polygon(poly_cutout, outline="red", width=2)
                cutout_pil.save(os.path.join(cutouts_outline_dir, cutout_fn))

                # Leaf-only cutout
                original_pil_image.crop((x1_p, y1_p, x2_p, y2_p)).save(os.path.join(cutouts_leaf_only_dir, cutout_fn))

                # Masked cutout
                leaf_mask = Image.new("L", (x2_p - x1_p, y2_p - y1_p), 0)
                ImageDraw.Draw(leaf_mask).polygon(poly_cutout, fill=255)
                masked_cutout = Image.new("RGB", (x2_p - x1_p, y2_p - y1_p), (0, 0, 0))
                masked_cutout.paste(original_pil_image.crop((x1_p, y1_p, x2_p, y2_p)), mask=leaf_mask)
                masked_cutout.save(os.path.join(cutouts_masked_dir, cutout_fn))

                # Binary mask
                binary_mask = Image.new("L", (x2 - x1, y2 - y1), 0)
                ImageDraw.Draw(binary_mask).polygon([(px - x1, py - y1) for px, py in polygon_pixels], fill=255)
                binary_mask.save(os.path.join(binary_masks_dir, cutout_fn))

        logging.info(f"  - Instances Detected: {total_detections}")
        logging.info(f"  - Instances Kept ({args.leaf_inclusion_mode}): {valid_instances}")
        
        annotated_frame_save_path = os.path.join(frame_dir, f"{os.path.splitext(filename)[0]}_annotated.png")
        annotated_image_for_frame.save(annotated_frame_save_path)

    csv_file.close()
    logging.info(f"Closed CSV file: {csv_filepath}")
    logging.info("Inference processing complete.")
    logging.info(f"Outputs saved in: {args.output_dir}")
    logging.info("--- Finished script: 23_inference_segmenter.py ---")

if __name__ == '__main__':
    main()