code/baselines/lhd/src/utils/utils.py

import numpy as np
import torch
from PIL import Image, ImageDraw
from typing import Optional, Union
import random
import pandas as pd
from PIL import Image
from scipy.ndimage import label, find_objects
from typing import Optional

def draw_obs(bg_img, mask):
    assert isinstance(bg_img, Image.Image)
    img = bg_img.copy()
    draw = ImageDraw.Draw(img)
    
    bg_width, bg_height = img.size 
    mask_height, mask_width = mask.shape

    x_scale = bg_width / mask_width
    y_scale = bg_height / mask_height

    for y in range(mask_height):
        for x in range(mask_width):
            if mask[y, x]: 
                scaled_x = round(x * x_scale)
                scaled_y = round(y * y_scale)
                draw.rectangle([(scaled_x, scaled_y), 
                                (scaled_x + round(x_scale), scaled_y + round(y_scale))], fill='black')
                
    return img

def draw_obstacles_pixel(bg, obstacle_masks):
    batch_size, _, _ = obstacle_masks.shape
    batched_images = []
    for i in range(batch_size):
        background = bg.copy()
        batched_images.append(draw_obs(background, obstacle_masks[i]))
    return batched_images     

def create_mask_from_binary_csv(csv_path, img_size, device):
    try:
        original_mask = pd.read_csv(csv_path, header=None).to_numpy()
    except FileNotFoundError:
        print(f"Error: The file {csv_path} was not found.")
        return None
    except Exception as e:
        print(f"Error reading or converting CSV file: {e}")
        return None

    mask_image = Image.fromarray((original_mask.astype(np.uint8) * 255), 'L')
    resized_image = mask_image.resize((img_size, img_size), Image.NEAREST)
    resized_mask_np = np.array(resized_image)
    final_mask = resized_mask_np / 255.0

    obstacle_mask_tensor = torch.tensor(
        final_mask, dtype=torch.float32, device=device
    ).unsqueeze(0)

    return obstacle_mask_tensor


def apply_random_obstacle_variations_fix_seed(
    base_obstacle_mask_tensor: torch.Tensor,
    num_variations: int,
    device: torch.device,
    seed: Optional[int] = None,
    scale_range: tuple = (0.9, 1.1),
    offset_range_pixels: int = 5,
    min_obstacle_area_ratio: float = 0.01,
    min_dist_between_obstacles: int = 0
) -> torch.Tensor:
    original_rng_state = None
    if seed is not None:
        original_rng_state = np.random.get_state()
        np.random.seed(seed)

    H, W = base_obstacle_mask_tensor.shape[1:]
    base_mask_np = base_obstacle_mask_tensor.squeeze(0).cpu().numpy().astype(np.uint8)

    labeled_array, num_features = label(base_mask_np)
    if num_features == 0:
        print("No obstacles found in the base mask. Returning original mask.")
        if seed is not None:
            np.random.set_state(original_rng_state)
        return base_obstacle_mask_tensor.repeat(num_variations, 1, 1)

    slices = find_objects(labeled_array)
    variant_masks_list = [base_mask_np]
    
    for _ in range(num_variations - 1):
        current_variant_mask = np.zeros_like(base_mask_np, dtype=np.uint8)
        
        for i, slc in enumerate(slices):
            if slc is None:
                continue
                
            y_slice, x_slice = slc
            obstacle_component = base_mask_np[y_slice, x_slice]
            
            scale = np.random.uniform(scale_range[0], scale_range[1])
            orig_h_comp, orig_w_comp = obstacle_component.shape
            new_h_comp, new_w_comp = int(orig_h_comp * scale), int(orig_w_comp * scale)
            
            if new_h_comp * new_w_comp < H * W * min_obstacle_area_ratio:
                new_h_comp = max(1, int(np.sqrt(H * W * min_obstacle_area_ratio)))
                new_w_comp = max(1, int(np.sqrt(H * W * min_obstacle_area_ratio)))

            pil_comp = Image.fromarray(obstacle_component * 255, 'L')
            scaled_comp_pil = pil_comp.resize((new_w_comp, new_h_comp), Image.NEAREST)
            scaled_comp = np.array(scaled_comp_pil) > 0

            center_y = y_slice.start + orig_h_comp // 2
            center_x = x_slice.start + orig_w_comp // 2
            offset_y = np.random.randint(-offset_range_pixels, offset_range_pixels + 1)
            offset_x = np.random.randint(-offset_range_pixels, offset_range_pixels + 1)
            new_center_y = center_y + offset_y
            new_center_x = center_x + offset_x

            new_y_start = new_center_y - new_h_comp // 2
            new_x_start = new_center_x - new_w_comp // 2

            new_y_start = max(0, min(new_y_start, H - new_h_comp))
            new_x_start = max(0, min(new_x_start, W - new_w_comp))
            new_y_end = new_y_start + new_h_comp
            new_x_end = new_x_start + new_w_comp

            current_variant_mask[new_y_start:new_y_end, new_x_start:new_x_end] = scaled_comp
        
        variant_masks_list.append(current_variant_mask)

    if seed is not None:
        np.random.set_state(original_rng_state)

    variant_masks_tensor = torch.tensor(
        np.stack(variant_masks_list), dtype=torch.float32, device=device
    )
    
    return variant_masks_tensor


def apply_shelf_obstacle_variations_fix_seed(
    base_obstacle_mask_tensor: torch.Tensor,
    num_variations: int,
    device: torch.device,
    seed: Optional[int] = None,
    scale_range: tuple = (0.9, 1.1),
    offset_range_pixels: int = 5,      
    min_obstacle_area_ratio: float = 0.01,
    num_special_variations: int = 0,    
    special_vertical_offset: int = 20   
) -> torch.Tensor:
    """
    Apply random variations to shelf obstacle mask with special variations
    
    Args:
        base_obstacle_mask_tensor: Base obstacle mask [1, H, W]
        num_variations: Number of variations to generate
        device: Torch device
        seed: Random seed for reproducibility
        scale_range: Range for scaling obstacles
        offset_range_pixels: Maximum pixel offset for obstacles
        min_obstacle_area_ratio: Minimum area ratio for obstacles
        num_special_variations: Number of special variations with vertical offset only
        special_vertical_offset: Vertical offset for special variations
        
    Returns:
        Varied obstacle masks [num_variations, 1, H, W]
    """
    original_rng_state = None
    if seed is not None:
        original_rng_state = np.random.get_state()
        np.random.seed(seed)

    H, W = base_obstacle_mask_tensor.shape[1:]
    base_mask_np = base_obstacle_mask_tensor.squeeze(0).cpu().numpy().astype(np.uint8)

    labeled_array, num_features = label(base_mask_np)
    if num_features == 0:
        print("No obstacles found in the base mask. Returning original mask.")
        if seed is not None:
            np.random.set_state(original_rng_state)
        return base_obstacle_mask_tensor.repeat(num_variations, 1, 1, 1)

    slices = find_objects(labeled_array) 
    variant_masks_list = []

    variant_masks_list.append(base_mask_np)
    
    for i in range(num_variations - 1):
        current_variation_index = i + 1 
        
        is_special = (current_variation_index >= (num_variations - num_special_variations))

        current_variant_mask = np.zeros_like(base_mask_np, dtype=np.uint8)
        transformed_obstacle_coords = [] 

        for j, slc in enumerate(slices):
            y_slice, x_slice = slc
            obstacle_component = base_mask_np[y_slice, x_slice]
            
            scale = np.random.uniform(scale_range[0], scale_range[1])
            orig_h_comp, orig_w_comp = obstacle_component.shape
            new_h_comp, new_w_comp = int(orig_h_comp * scale), int(orig_w_comp * scale)
            
            if new_h_comp * new_w_comp < H * W * min_obstacle_area_ratio:
                new_h_comp = max(1, int(np.sqrt(H*W*min_obstacle_area_ratio)))
                new_w_comp = max(1, int(np.sqrt(H*W*min_obstacle_area_ratio)))

            pil_comp = Image.fromarray(obstacle_component * 255, 'L')
            scaled_comp_pil = pil_comp.resize((new_w_comp, new_h_comp), Image.NEAREST)
            scaled_comp = np.array(scaled_comp_pil) > 0 

            center_y = y_slice.start + orig_h_comp // 2
            center_x = x_slice.start + orig_w_comp // 2

            if is_special:
                offset_x = 0
                offset_y = np.random.choice([-special_vertical_offset, special_vertical_offset])
            else:
                offset_y = np.random.randint(-offset_range_pixels, offset_range_pixels + 1)
                offset_x = np.random.randint(-offset_range_pixels, offset_range_pixels + 1)

            new_center_y = center_y + offset_y
            new_center_x = center_x + offset_x

            new_y_start = new_center_y - new_h_comp // 2
            new_x_start = new_center_x - new_w_comp // 2

            new_y_start = max(0, min(new_y_start, H - new_h_comp))
            new_x_start = max(0, min(new_x_start, W - new_w_comp))
            new_y_end = new_y_start + new_h_comp
            new_x_end = new_x_start + new_w_comp

            for prev_y_slice, prev_x_slice in transformed_obstacle_coords:
                if (max(new_y_start, prev_y_slice.start) < min(new_y_end, prev_y_slice.stop) and
                    max(new_x_start, prev_x_slice.start) < min(new_x_end, prev_x_slice.stop)):
                    pass 
            
            current_variant_mask[new_y_start:new_y_end, new_x_start:new_x_end] |= scaled_comp

            transformed_obstacle_coords.append(
                (slice(new_y_start, new_y_end), slice(new_x_start, new_x_end))
            )
        
        variant_masks_list.append(current_variant_mask)
    
    final_variant_masks = torch.tensor(
        np.stack(variant_masks_list), dtype=torch.float32, device=device
    ).unsqueeze(1) 

    if seed is not None:
        np.random.set_state(original_rng_state)

    return final_variant_masks


def apply_room_obstacle_variations_fix_seed( 
    horizontal_mask_tensor: torch.Tensor, 
    vertical_mask_tensor: torch.Tensor,   
    num_variations: int,
    device: torch.device,
    seed: Optional[int] = None,
    scale_range: tuple = (0.9, 1.1),
    offset_range_pixels: int = 5,
    min_obstacle_area_ratio: float = 0.01 
) -> torch.Tensor:

    original_rng_state = None
    if seed is not None:
        original_rng_state = np.random.get_state()
        np.random.seed(seed)

    H, W = horizontal_mask_tensor.shape[1:]
    h_mask_np = horizontal_mask_tensor.squeeze(0).cpu().numpy().astype(np.uint8)
    v_mask_np = vertical_mask_tensor.squeeze(0).cpu().numpy().astype(np.uint8)

    variant_masks_list = []

    original_plus_mask = np.logical_or(h_mask_np, v_mask_np).astype(np.uint8)
    variant_masks_list.append(original_plus_mask)

    for _ in range(num_variations - 1):
        
        is_horizontal = False
        
        if np.random.rand() < 0.5:
            mask_to_transform = h_mask_np
            static_mask = v_mask_np
            is_horizontal = True 
        else:
            mask_to_transform = v_mask_np
            static_mask = h_mask_np
            is_horizontal = False 

        labeled_array, num_features = label(mask_to_transform)
        if num_features == 0:
            variant_masks_list.append(static_mask)
            continue
            
        slices = find_objects(labeled_array)
        slc = slices[0] 
        y_slice, x_slice = slc
        obstacle_component = mask_to_transform[y_slice, x_slice]
            
        scale = np.random.uniform(scale_range[0], scale_range[1])
        orig_h_comp, orig_w_comp = obstacle_component.shape
        new_h_comp, new_w_comp = int(orig_h_comp * scale), int(orig_w_comp * scale)
        
        if new_h_comp * new_w_comp < H * W * min_obstacle_area_ratio:
            new_h_comp = max(1, int(np.sqrt(H*W*min_obstacle_area_ratio)))
            new_w_comp = max(1, int(np.sqrt(H*W*min_obstacle_area_ratio)))

        pil_comp = Image.fromarray(obstacle_component * 255, 'L')
        scaled_comp_pil = pil_comp.resize((new_w_comp, new_h_comp), Image.NEAREST)
        scaled_comp = np.array(scaled_comp_pil) > 0 

        center_y = y_slice.start + orig_h_comp // 2
        center_x = x_slice.start + orig_w_comp // 2
        
        offset_y = 0
        offset_x = 0
        if is_horizontal:
            offset_y = np.random.randint(-offset_range_pixels, offset_range_pixels + 1)
        else:
            offset_x = np.random.randint(-offset_range_pixels, offset_range_pixels + 1)

        new_center_y = center_y + offset_y
        new_center_x = center_x + offset_x

        new_y_start = new_center_y - new_h_comp // 2
        new_x_start = new_center_x - new_w_comp // 2
        new_y_start = max(0, min(new_y_start, H - new_h_comp))
        new_x_start = max(0, min(new_x_start, W - new_w_comp))
        new_y_end = new_y_start + new_h_comp
        new_x_end = new_x_start + new_w_comp
        
        varied_bar_mask = np.zeros_like(h_mask_np, dtype=np.uint8)
        varied_bar_mask[new_y_start:new_y_end, new_x_start:new_x_end] |= scaled_comp

        final_mask = np.logical_or(varied_bar_mask, static_mask).astype(np.uint8)
        variant_masks_list.append(final_mask)
    
    final_variant_masks = torch.tensor(
        np.stack(variant_masks_list), dtype=torch.float32, device=device
    ).unsqueeze(1)

    if seed is not None:
        np.random.set_state(original_rng_state)

    return final_variant_masks


def is_valid_obstacles(mask, x_start, x_size, y_start, y_size, min_distance):
    x_end, y_end = x_start + x_size, y_start + y_size
    
    x_safe_start = max(0, x_start - min_distance)
    y_safe_start = max(0, y_start - min_distance)
    x_safe_end = min(mask.shape[0], x_end + min_distance)
    y_safe_end = min(mask.shape[1], y_end + min_distance)
    
    if torch.sum(mask[x_safe_start:x_safe_end, y_safe_start:y_safe_end]) > 0:
        return False
    return True


def randgen_obstacle_mask(batch_size, image_size, seed:Optional[int]=None, device='cuda'):
    if seed is not None:
        torch.manual_seed(seed)
        random.seed(seed)
    
    min_area_ratio = 0.07
    max_area_ratio = 0.3
    random_ratio = random.uniform(min_area_ratio, max_area_ratio)
    
    total_obstacle_area = int(image_size**2 * random_ratio) 
    max_length = int(image_size*5/6) 
    min_length = int(image_size*2/15) 
    min_distance = int(image_size/15)
    max_attempts = 50

    masks = torch.zeros(batch_size, image_size, image_size, dtype=torch.bool, device=device)

    for i in range(batch_size):
        num_obstacles = torch.randint(2, 5, (1,)).item()
        areas = torch.tensor(total_obstacle_area // num_obstacles, device=device).repeat(num_obstacles)

        for area in areas:
            attempts = 0
            while attempts < max_attempts:
                possible_x_sizes = [x for x in range(min_length, max_length + 1) if min_length <= (area // x) <= max_length]
                if not possible_x_sizes:
                    area = (area * 0.8).long()
                    attempts += 1
                    continue

                x_size = random.choice(possible_x_sizes)
                y_size = area // x_size

                x_start = torch.randint(0, image_size - x_size + 1, (1,)).item()
                y_start = torch.randint(0, image_size - y_size + 1, (1,)).item()

                if is_valid_obstacles(masks[i], x_start, x_size, y_start, y_size, min_distance):
                    masks[i, x_start:x_start + x_size, y_start:y_start + y_size] = 1
                    break
                attempts += 1
    return masks


def get_distance(a, b):
    return torch.sqrt((a[0] - b[0])**2 + (a[1] - b[1])**2)


def is_valid_goals(pixel_x, pixel_y, obstacles, batch_idx, clearance=0):
    if obstacles.dim() == 4:
        obstacles = obstacles.squeeze(1) 
    
    B, H, W = obstacles.shape
    
    if not (0 <= pixel_x < W and 0 <= pixel_y < H):
        return False

    x_start = max(0, pixel_x - clearance)
    x_end   = min(W, pixel_x + clearance + 1) 
    y_start = max(0, pixel_y - clearance)
    y_end   = min(H, pixel_y + clearance + 1)

    area_to_check = obstacles[batch_idx, x_start:x_end, y_start:y_end]
    
    return torch.sum(area_to_check) == 0


def gen_goals(
    bounds, 
    n: Union[tuple, int], 
    img_size: int,
    dist: Optional[float] = 0.25, 
    obstacles: Optional[torch.Tensor] = None,
    seed: Optional[int] = None,
    device: str = 'cuda',
    obstacle_clearance_pixels: int = 0
):
    if seed is not None:
        torch.manual_seed(seed)
    
    assert len(bounds) == 4, f'Inappropriate map bound: {bounds}'
    if isinstance(n, int):
        M = 1
        N = n
    else:
        M, N = n
        
    goals = []
    for batch_idx in range(N):
        batch_goals = []
        while len(batch_goals) < M:
            x = bounds[0] + (bounds[1] - bounds[0]) * torch.rand((1, 1), dtype=torch.float32, device=device)
            y = bounds[2] + (bounds[3] - bounds[2]) * torch.rand((1, 1), dtype=torch.float32, device=device)

            if obstacles is not None:
                pixel_x = ((x + 1) * 0.5 * (img_size - 1)).long().item()
                pixel_y = ((y + 1) * 0.5 * (img_size - 1)).long().item()

                if not is_valid_goals(pixel_x, pixel_y, obstacles, batch_idx, clearance=obstacle_clearance_pixels):
                    continue 

            valid = True
            if dist is not None:
                for existing_goal in batch_goals:
                    if get_distance([x, y], [existing_goal[0, 0], existing_goal[0, 1]]) < dist:
                        valid = False
                        break
            
            if valid:
                batch_goals.append(torch.cat((x, y), dim=1))

        goals.append(torch.cat(batch_goals, dim=0))
                    
    return torch.stack(goals)


def overlay_goals(img, img_size, objs, pos, goal_size):
    assert pos.dim() == 3
    n = len(objs)
    
    if len(img) != pos.shape[0]:
        img = img * pos.shape[0]
    if img[0].height != img_size:
        new_size = (img_size, img_size)
        for i in range(len(img)):
            img[i] = img[i].resize(new_size, Image.LANCZOS)
    
    W, H = img[0].size
    
    objs = [img.copy() for img in objs]
    for i in range(len(objs)):
        objs[i] = objs[i].resize((W // goal_size, H // goal_size), Image.LANCZOS)
    pos_pix = ((1 + pos) / 2 * torch.tensor([H - 1, W - 1], device=pos.device, dtype=pos.dtype))
    
    imgs = []
    for i, ps in enumerate(pos_pix.cpu().numpy()):
        bg = img[i].copy()
        for j, p in enumerate(ps):
            p1, p0 = round(p[1]), round(p[0])
            w, h = objs[j].size
            bg.paste(objs[j], (p1 - w // 2, p0 - h // 2), objs[j])
        img_np = np.array(bg)[..., :3] / 255
        imgs.append(img_np)
    
    return torch.tensor(imgs, dtype=pos.dtype, device=pos.device).permute(0, 3, 1, 2)