import gradio as gr
import numpy as np
import torch
import cv2
from PIL import Image
import os
from dataclasses import dataclass
from typing import List, Dict, Any, Optional
from tqdm import tqdm
import time
import platform
import multiprocessing
from multiprocessing import Process, Queue, Value, Lock, Manager
import queue
from multiprocessing import Value
# Assuming these are available in your environment
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
from GroundingDINO.groundingdino.util.inference import load_model, load_image, predict
from torchvision.ops import box_convert

# --- Configuration (Adjust paths as needed) ---
SAM2_CHECKPOINT = os.path.abspath("models/sam2.1_hiera_large.pt")
SAM2_MODEL_CONFIG = os.path.abspath("models/sam2.1_hiera_l.yaml")
GROUNDING_DINO_CONFIG = os.path.abspath("models/GroundingDINO_SwinT_OGC.py")
GROUNDING_DINO_CHECKPOINT = os.path.abspath("models/groundingdino_swint_ogc.pth")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
OUTPUT_DIR = "outputs"
MASKS_DIR = "generated_masks"

# --- Make input_folder truly global ---
input_folder = ""  # Initialize globally


# --- Synchronization ---
debug = False
stop_processing = Value('b', False)  # Initialize as a multiprocessing Value

def open_folder():
    open_folder_path = os.path.abspath("outputs")
    if platform.system() == "Windows":
        os.startfile(open_folder_path)
    elif platform.system() == "Linux":
        os.system(f'xdg-open "{open_folder_path}"')

def open_folder_mask():
    open_folder_path = os.path.abspath("generated_masks")
    if platform.system() == "Windows":
        os.startfile(open_folder_path)
    elif platform.system() == "Linux":
        os.system(f'xdg-open "{open_folder_path}"')

# --- Model Loading (For Single Image Processing) ---
# Check if files exist
if not os.path.exists(SAM2_CHECKPOINT):
    print(f"Error: SAM2 checkpoint file not found at {SAM2_CHECKPOINT}")
    exit()
if not os.path.exists(SAM2_MODEL_CONFIG):
    print(f"Error: SAM2 model config file not found at {SAM2_MODEL_CONFIG}")
    exit()
if not os.path.exists(GROUNDING_DINO_CONFIG):
    print(f"Error: Grounding DINO config file not found at {GROUNDING_DINO_CONFIG}")
    exit()
if not os.path.exists(GROUNDING_DINO_CHECKPOINT):
    print(f"Error: Grounding DINO checkpoint file not found at {GROUNDING_DINO_CHECKPOINT}")
    exit()

# --- Data Classes ---

@dataclass
class BoundingBox:
    xmin: int
    ymin: int
    ymax: int
    xmax: int

    @property
    def xyxy(self) -> List[float]:
        return [self.xmin, self.ymin, self.xmax, self.ymax]

@dataclass
class DetectionResult:
    score: float
    label: str
    box: BoundingBox
    mask: Optional[np.array] = None

# --- Helper Functions ---

def mask_to_polygon(mask: np.ndarray) -> List[List[int]]:
    mask = (mask > 0).astype(np.uint8)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return []
    largest_contour = max(contours, key=cv2.contourArea)
    epsilon = 0.001 * cv2.arcLength(largest_contour, True)
    approx_contour = cv2.approxPolyDP(largest_contour, epsilon, True)
    return approx_contour.reshape(-1, 2).tolist()

def polygon_to_mask(polygon: List, image_shape: tuple) -> np.ndarray:
    if not polygon:
        return np.zeros(image_shape, dtype=np.uint8)
    mask = np.zeros(image_shape, dtype=np.uint8)
    pts = np.array(polygon, dtype=np.int32)
    cv2.fillPoly(mask, [pts], color=(1,))
    return mask

def get_boxes_from_gdino(results: tuple) -> List[List[float]]:
    boxes = []
    for box in results[0]:
        xyxy = box.tolist()
        boxes.append(xyxy)
    return boxes
    
def refine_masks(masks: np.ndarray, image_shape: tuple) -> List[np.ndarray]:
    refined_masks = []
    for mask in masks:
        # Ensure the mask is a boolean array
        mask_bool = mask > 0

        # Resize the mask to the original image shape
        mask_resized = cv2.resize(mask_bool.astype(np.uint8) * 255, (image_shape[1], image_shape[0]))

        # Threshold the resized mask to ensure binary values (0 or 255)
        _, mask_thresholded = cv2.threshold(mask_resized, 127, 255, cv2.THRESH_BINARY)

        # Convert the mask back to boolean array (True/False)
        mask_bool = mask_thresholded > 0
        
        refined_masks.append(mask_bool)

    return refined_masks

# --- Main Processing Functions ---
def detect(
    temp_image_path: str,
    labels: List[str],
    box_threshold: float = 0.3,
    text_threshold: float = 0.25,
    local_grounding_model = None
) -> List[DetectionResult]:
    image_source, image = load_image(temp_image_path)
    h, w, _ = image_source.shape
    
    text_prompt = ". ".join(labels) + "."
    
    with torch.amp.autocast(device_type='cuda', dtype=torch.float16): # Use updated autocast
        boxes, confidences, phrases = predict(
            model= local_grounding_model if local_grounding_model else grounding_model,
            image=image,
            caption=text_prompt,
            box_threshold=box_threshold,
            text_threshold=text_threshold,
        )
    
    # Scale boxes by image dimensions
    boxes = boxes * torch.Tensor([w, h, w, h])
    boxes = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
    
    detections = []
    for box, score, label in zip(boxes, confidences, phrases):
        detections.append(
            DetectionResult(
                score=score,
                label=label,
                box=BoundingBox(
                    xmin=int(box[0]),
                    ymin=int(box[1]),
                    xmax=int(box[2]),
                    ymax=int(box[3])
                )
            )
        )
    return detections

def segment(
    temp_image_path: str,
    detection_results: List[DetectionResult],
    local_sam2_predictor = None
) -> List[DetectionResult]:
    image_source, _ = load_image(temp_image_path)
    local_sam2_predictor.set_image(image_source)

    # Get boxes in correct format
    boxes = np.array([det.box.xyxy for det in detection_results])
    
    # Use SAM2 predictor
    masks, _, _ = local_sam2_predictor.predict(
        point_coords=None,
        point_labels=None,
        box=boxes,
        multimask_output=False,
    )
    
    # Handle mask dimensionality
    if masks.ndim == 4:
        masks = masks.squeeze(1)
    
    # Assign masks to detection results
    for detection_result, mask in zip(detection_results, masks):
        detection_result.mask = mask > 0  # Convert to boolean mask
        
    return detection_results
    
def process_image(input_image, text_prompt, box_threshold, text_threshold, save_mask, local_sam2_predictor, local_grounding_model, output_dir, masks_dir, processed_files_queue, processed_count, lock, image_path=None):
    global input_folder

    try:
        # If in batch mode, input_image is a path string
        if processed_files_queue is not None:
            # Batch mode: Use image_path directly
            temp_image_path = image_path
            input_image = Image.open(image_path).convert("RGB")  # Read the image here
        else:
            # Single image mode: input_image is a numpy array
            input_image = Image.fromarray(input_image)

            # Make temp file unique per process
            pid = multiprocessing.current_process().pid
            timestamp = int(time.time())
            temp_filename = f"temp_image_{timestamp}_{pid}.png"

            # Use os.path.join for paths
            temp_image_path = os.path.join(OUTPUT_DIR, temp_filename)

            # Create directories if they don't exist
            os.makedirs(os.path.dirname(temp_image_path), exist_ok=True)
            input_image.save(temp_image_path)

        # Detect and segment
        labels = [text_prompt]
        detections = detect(temp_image_path, labels, box_threshold, text_threshold, local_grounding_model=local_grounding_model)

        if not detections:
            print(f"No detections found for image with prompt '{text_prompt}'")
            return np.array(input_image), None, [], temp_image_path

        detections = segment(temp_image_path, detections, local_sam2_predictor=local_sam2_predictor)

        # Create visualization (this is the segmented image)
        segmented_image = annotate(input_image, detections)

        # Save the segmented image with a unique name
        if processed_files_queue is None:
            output_image_filename = f"output_{timestamp}_{text_prompt.replace(' ', '_')}_{box_threshold:.2f}_{pid}.png"
            output_image_path = os.path.join(output_dir, output_image_filename)
            Image.fromarray(segmented_image).save(output_image_path) # Save segmented_image

        # Handle mask generation and coordinates
        mask_image = None
        mask_coordinates = []

        if save_mask and detections:
            # Create binary mask
            combined_mask = np.zeros(np.array(input_image).shape[:2], dtype=np.uint8)
            for detection in detections:
                if detection.mask is not None:
                    combined_mask |= detection.mask.astype(np.uint8)

            # Convert to image
            mask_image = Image.fromarray(combined_mask * 255)

            # Get coordinates
            mask_coords = np.where(combined_mask > 0)
            mask_coordinates = list(zip(mask_coords[0].tolist(), mask_coords[1].tolist()))

            # Save mask **only if not in batch mode**
            if processed_files_queue is None:
                os.makedirs(masks_dir, exist_ok=True)
                mask_filename = f"mask_{timestamp}_{text_prompt.replace(' ', '_')}_{box_threshold:.2f}_{pid}.png"
                mask_image.save(os.path.join(masks_dir, mask_filename))

        # Cleanup: only remove temp files made by this process in single-image mode
        if processed_files_queue is None and str(pid) in temp_image_path:
            os.remove(temp_image_path)

        return segmented_image, mask_image, mask_coordinates, temp_image_path

    except Exception as e:
        print(f"Error processing image: {e}")
        return np.array(input_image), None, [], temp_image_path

def annotate(image: Image.Image, detection_results: List[DetectionResult]) -> np.ndarray:
    image_cv2 = np.array(image)
    image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)

    for detection in detection_results:
        label = detection.label
        score = detection.score
        box = detection.box
        mask = detection.mask

        color = np.random.randint(0, 256, size=3)

        if mask is not None:
            overlay = image_cv2.copy()
            mask_display = (mask * 255).astype(np.uint8)
            overlay[mask == 1] = color
            image_cv2 = cv2.addWeighted(image_cv2, 0.7, overlay, 0.3, 0)

            contours, _ = cv2.findContours(mask_display, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(image_cv2, contours, -1, color.tolist(), 2)

        cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color.tolist(), 2)
        cv2.putText(
            image_cv2,
            f'{label}: {score:.2f}',
            (box.xmin, box.ymin - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            color.tolist(),
            2
        )

    return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)

def worker_process(
    gpu_id: int,
    task_queue: Queue,
    results_queue: Queue,
    output_segmented_folder: str,
    output_masks_folder: str,
    text_prompt: str,
    box_threshold: float,
    text_threshold: float,
    batch_size: int,
    processed_files_queue: Queue,
    processed_count: Value,
    big_lock: Lock,
    file_locks: dict,
    small_lock: Lock,
    start_times: dict,
    stop_processing: Value
):
    try:
        # Record start time for this worker
        start_times[gpu_id] = time.time()

        # Set device for torch
        device = f"cuda:{gpu_id}"
        print(f"Worker {gpu_id} starting on device {device}")

        # Load models
        local_sam2_model = build_sam2(SAM2_MODEL_CONFIG, SAM2_CHECKPOINT, device=device)
        local_sam2_predictor = SAM2ImagePredictor(local_sam2_model)
        local_grounding_model = load_model(
            model_config_path=GROUNDING_DINO_CONFIG,
            model_checkpoint_path=GROUNDING_DINO_CHECKPOINT,
            device=device
        )

        while not stop_processing.value:
            # DEBUG: Print when a worker is waiting for the big_lock
            if debug:
                print(f"Worker {gpu_id}: Waiting for big_lock...")

            # Acquire the "big lock" before getting a task
            big_lock.acquire()

            # DEBUG: Print when a worker has acquired the big_lock
            if debug:
                print(f"Worker {gpu_id}: Acquired big_lock")

            try:
                image_path = task_queue.get_nowait()
                if image_path == "DONE":
                    if debug:
                        print(f"Worker {gpu_id} received DONE signal")
                    task_queue.put("DONE")  # Put it back for other workers
                    big_lock.release()
                    if debug:
                        print(f"Worker {gpu_id}: Released big_lock (DONE)")
                    break

                # Generate output file paths
                base_name = os.path.splitext(os.path.basename(image_path))[0]
                segmented_image_path = os.path.join(output_segmented_folder, f"segmented_{base_name}.png")
                mask_image_path = os.path.join(output_masks_folder, f"mask_{base_name}.png")

                # DEBUG: Print the file_locks content before checking
                if debug:
                    print(f"Worker {gpu_id}: file_locks content: {file_locks}")

                # Immediately mark files as "in progress"
                if segmented_image_path in file_locks or mask_image_path in file_locks:
                    results_queue.put(f"Skipped {base_name} - already being processed by another worker")

                    # DEBUG: Print which worker is skipping the file
                    if debug:
                        print(f"Worker {gpu_id}: Skipping {base_name} (already in file_locks)")

                    big_lock.release()
                    if debug:
                        print(f"Worker {gpu_id}: Released big_lock (skip)")
                    continue
                else:
                    file_locks[segmented_image_path] = True  # Mark as being processed
                    file_locks[mask_image_path] = True

                    # DEBUG: Print which worker is marking the file as in progress
                    if debug:
                        print(f"Worker {gpu_id}: Marked {base_name} as in progress")

                # Release the "big lock" after marking files
                big_lock.release()
                if debug:
                    print(f"Worker {gpu_id}: Released big_lock (processing)")

                # DEBUG: Print before starting processing
                if debug:
                    print(f"Worker {gpu_id}: Starting processing of {base_name}...")

                # Process the image
                input_image = Image.open(image_path).convert("RGB")

                # Pass arguments to process_image in the correct order:
                segmented_image, mask_image, _, _ = process_image(
                    np.array(input_image),
                    text_prompt,
                    box_threshold,
                    text_threshold,
                    save_mask=True,
                    local_sam2_predictor=local_sam2_predictor,
                    local_grounding_model=local_grounding_model,
                    output_dir=output_segmented_folder,
                    masks_dir=output_masks_folder,
                    processed_files_queue=processed_files_queue,
                    processed_count=processed_count,
                    lock=small_lock,
                    image_path=image_path
                )

                # Save outputs
                Image.fromarray(segmented_image).save(segmented_image_path)
                if mask_image:
                    mask_image.save(mask_image_path)

                # Update processed files and count
                with small_lock:
                    processed_files_queue.put(segmented_image_path)
                    processed_count.value += 1
                    # DEBUG: Print file_locks before releasing the lock
                    if debug:
                        print(f"Worker {gpu_id}: file_locks before releasing small_lock: {file_locks}")
                    del file_locks[segmented_image_path]
                    del file_locks[mask_image_path]
                    if debug:
                        print(f"Worker {gpu_id}: Finished processing {base_name}, removed from file_locks")

                results_queue.put(f"Processed {base_name}")

            except queue.Empty:
                big_lock.release()
                if debug:
                    print(f"Worker {gpu_id}: Released big_lock (empty queue)")
                break
            except Exception as e:
                results_queue.put(f"Error processing {image_path}: {str(e)}")
                try:
                    # Just try to release the lock; it handles the unlocked case
                    big_lock.release()
                except:
                    pass
                if debug:    
                    print(f"Worker {gpu_id}: Released big_lock (exception)")

    except Exception as e:
        results_queue.put(f"Worker {gpu_id} failed: {str(e)}")
    finally:
        if debug:
            print(f"Worker {gpu_id} shutting down")

def batch_process_images(
    input_folder_local,
    output_segmented_folder,
    output_masks_folder,
    text_prompt,
    box_threshold,
    text_threshold,
    batch_size,
    gpu_ids,
):
    global input_folder, stop_processing
    input_folder = input_folder_local

    try:
        # Create output directories
        os.makedirs(output_segmented_folder, exist_ok=True)
        os.makedirs(output_masks_folder, exist_ok=True)

        # Get list of image files
        image_files = [
            os.path.join(input_folder, f)
            for f in os.listdir(input_folder)
            if os.path.isfile(os.path.join(input_folder, f))
            and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp'))
        ]
        image_files = list(set(image_files))
        total_tasks = len(image_files)

        if total_tasks == 0:
            return "No images found in input folder"

        print(f"Found {total_tasks} images to process")

        # Create a new Manager for each batch processing run
        manager = Manager()
        try:
            # Create shared objects using the manager
            task_queue = manager.Queue()
            results_queue = manager.Queue()
            processed_files_queue = manager.Queue()
            processed_count = manager.Value('i', 0)
            big_lock = manager.Lock()
            file_locks = manager.dict()
            small_lock = manager.Lock()
            start_times = manager.dict()
            stop_processing = manager.Value('b', False)

            # Fill task queue
            for image_path in image_files:
                task_queue.put(image_path)
            # Add termination signals
            for _ in range(len(gpu_ids)):
                task_queue.put("DONE")

            # Start worker processes
            processes = []
            for gpu_id in gpu_ids:
                p = Process(
                    target=worker_process,
                    args=(
                        gpu_id,
                        task_queue,
                        results_queue,
                        output_segmented_folder,
                        output_masks_folder,
                        text_prompt,
                        box_threshold,
                        text_threshold,
                        batch_size,
                        processed_files_queue,
                        processed_count,
                        big_lock,
                        file_locks,
                        small_lock,
                        start_times,
                        stop_processing
                    )
                )
                p.start()
                processes.append(p)

            # Monitor progress and calculate statistics
            start_time = time.time()
            last_processed_count = 0
            while processed_count.value < total_tasks and any(p.is_alive() for p in processes):
                try:
                    result = results_queue.get(timeout=60)
                    print(result)

                    # Calculate and display statistics
                    elapsed_time = time.time() - start_time
                    current_processed_count = processed_count.value
                    processing_rate = current_processed_count / elapsed_time if elapsed_time > 0 else 0
                    remaining_tasks = total_tasks - current_processed_count
                    estimated_time_remaining = remaining_tasks / processing_rate if processing_rate > 0 else 0
                    
                    print(f"Processed: {current_processed_count}/{total_tasks} | "
                          f"Speed: {processing_rate:.2f} images/s | "
                          f"ETA: {estimated_time_remaining:.2f} s")
                    
                    last_processed_count = current_processed_count

                except queue.Empty:
                    print("No progress in the last minute, checking if workers are alive...")
                    if not any(p.is_alive() for p in processes):
                        print("All workers have died")
                        break

            # Wait for all processes to finish
            for p in processes:
                p.join(timeout=5)
                if p.is_alive():
                    p.terminate()

            return f"Completed processing {processed_count.value} out of {total_tasks} images"

        finally:
            # Clean up manager resources
            manager.shutdown()

    except Exception as e:
        return f"Batch processing failed: {str(e)}"

# Update the Gradio interface to handle the new return value
def create_gradio_interface():
    with gr.Blocks() as interface:
        with gr.Tab("Single Image Processing"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="numpy", label="Input Image", height=500)
                    text_prompt = gr.Textbox(
                        label="Enter object to segment. Make sure it is lowercase and each word ends with . Also plural or singular makes difference",
                        value="eye. lips. ear.",
                        placeholder="Enter object to segment. Make sure it is lowercase and each word ends with . ",
                    )
                    box_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.35,
                        label="Box Threshold",
                    )
                    text_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.25,
                        label="Text Threshold",
                    )
                    run_button = gr.Button("Process Image")
                with gr.Column():
                    output_image = gr.Image(type="numpy", label="Output Image", height=500)
                    btn_open_outputs = gr.Button("Open Outputs Folder")
                    btn_open_outputs.click(fn=open_folder)
                    output_mask = gr.Image(type="numpy", label="Output Mask", height=500)
                    btn_open_outputs = gr.Button("Open Generated Masks Folder")
                    btn_open_outputs.click(fn=open_folder_mask)
                    mask_coordinates_output = gr.Textbox(label="Mask Coordinates", visible=False)

            def process_and_return_visible(input_image, prompt, box_threshold, text_threshold, save_mask):
                # Initialize models and other resources
                local_sam2_model = build_sam2(SAM2_MODEL_CONFIG, SAM2_CHECKPOINT, device=DEVICE)
                local_sam2_predictor = SAM2ImagePredictor(local_sam2_model)
                local_grounding_model = load_model(
                    model_config_path=GROUNDING_DINO_CONFIG,
                    model_checkpoint_path=GROUNDING_DINO_CHECKPOINT,
                    device=DEVICE
                )

                output_img, mask_img, mask_coords, _ = process_image(
                    input_image, prompt, box_threshold, text_threshold, save_mask,
                    local_sam2_predictor, local_grounding_model, OUTPUT_DIR, MASKS_DIR,
                    None, None, None  # Pass None for unused parameters in single image processing
                )
                mask_coords_str = str(mask_coords)
                return output_img, mask_img, mask_coords_str

            run_button.click(
                fn=process_and_return_visible,
                inputs=[input_image, text_prompt, box_threshold, text_threshold, gr.Checkbox(True, label="Save Mask")],
                outputs=[output_image, output_mask, mask_coordinates_output],
            )

        with gr.Tab("Batch Processing"):
            with gr.Row():
                with gr.Column():
                    # Use a gr.Textbox to get input_folder value
                    input_folder_textbox = gr.Textbox(label="Input Folder", placeholder="Enter path to input folder")
                    output_segmented_folder = gr.Textbox(label="Output Segmented Images Folder", placeholder="Enter path to output segmented images folder")
                    output_masks_folder = gr.Textbox(label="Output Masks Folder", placeholder="Enter path to output masks folder")
                    batch_text_prompt = gr.Textbox(
                        label="Enter object to segment. Make sure it is lowercase and each word ends with . Also plural or singular makes difference",
                        value="eye. lips. ear.",
                        placeholder="Enter object to segment (e.g., 'eyes', 'face')",
                    )
                    batch_box_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.35,
                        label="Box Threshold",
                    )
                    batch_text_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.25,
                        label="Text Threshold",
                    )
                    batch_size_input = gr.Number(label="Batch Size", value=4)
                    gpu_ids_input = gr.Textbox(
                        label="GPU IDs (comma-separated)",
                        placeholder="e.g., 0,1",
                        value="0"
                    )
                    batch_run_button = gr.Button("Batch Process Images")
                    stop_batch_button = gr.Button("Stop Batch Processing")
                with gr.Column():
                    batch_status = gr.Textbox(label="Status", interactive=False)

            # Connect batch_run_button.click to an intermediate function
            def intermediate_batch_process(input_folder_local, output_segmented_folder, output_masks_folder, text_prompt, box_threshold, text_threshold, batch_size, gpu_ids_str):
                global input_folder, stop_processing
                input_folder = input_folder_local
                
                try:
                    # Reset multiprocessing resources
                    if 'stop_processing' in globals():
                        del stop_processing
                    
                    # Create new stop_processing Value
                    stop_processing = Value('b', False)

                    return batch_process_images(
                        input_folder_local,
                        output_segmented_folder,
                        output_masks_folder,
                        text_prompt,
                        box_threshold,
                        text_threshold,
                        batch_size=int(batch_size),
                        gpu_ids=[int(x.strip()) for x in gpu_ids_str.split(',')]
                    ) or "Batch processing completed"
                except Exception as e:
                    return f"Error in batch processing: {str(e)}"

            def stop_batch_processing():
                global stop_processing
                try:
                    if stop_processing is None:
                        stop_processing = Value('b', True)
                    else:
                        stop_processing.value = True
                    return "Stopping batch processing..."
                except Exception as e:
                    return f"Error while stopping batch processing: {str(e)}"

            batch_run_button.click(
                fn=intermediate_batch_process,
                inputs=[
                    input_folder_textbox,
                    output_segmented_folder,
                    output_masks_folder,
                    batch_text_prompt,
                    batch_box_threshold,
                    batch_text_threshold,
                    batch_size_input,
                    gpu_ids_input
                ],
                outputs=batch_status
            )

            stop_batch_button.click(
                fn=stop_batch_processing,
                outputs=batch_status
            )

    return interface

if __name__ == "__main__":
    try:
        # Try to set start method
        multiprocessing.set_start_method('spawn')
    except RuntimeError:
        # If already set, get a spawn context
        ctx = multiprocessing.get_context('spawn')
        # Replace Process usage with ctx.Process throughout the code
        globals()['Process'] = ctx.Process
        globals()['Queue'] = ctx.Queue
        globals()['Value'] = ctx.Value
        globals()['Lock'] = ctx.Lock
        globals()['Manager'] = ctx.Manager
    
    demo = create_gradio_interface()
    demo.launch(share=False, inbrowser=True)