utils.py

import cv2
from typing import List
import numpy as np
from PIL import Image
import re
import random


def save_array_to_img(img_arr, img_p):
    Image.fromarray(img_arr.astype(np.uint8)).save(img_p)

def dilate_mask(mask, dilate_factor=15):
    mask = mask.astype(np.uint8)
    mask = cv2.dilate(
        mask,
        np.ones((dilate_factor, dilate_factor), np.uint8),
        iterations=1
    )
    return mask

def show_points(ax, coords: List[List[float]], labels: List[int], size=375):
    coords = np.array(coords)
    labels = np.array(labels)
    color_table = {0: 'red', 1: 'green'}
    for label_value, color in color_table.items():
        points = coords[labels == label_value]
        ax.scatter(points[:, 0], points[:, 1], color=color, marker='*',
                   s=size, edgecolor='white', linewidth=1.25)
        
def show_mask(ax, mask: np.ndarray, random_color=False):
    mask = mask.astype(np.uint8)
    if np.max(mask) == 255:
        mask = mask / 255
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
    h, w = mask.shape[-2:]
    mask_img = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_img)
    
def get_centers(binary_mask):
    
    num_labels, labels_im = cv2.connectedComponents(binary_mask.astype(np.uint8))
    if num_labels <= 1:
        return None  

    all_coords = []
    centers = []
    
    for label in range(1, num_labels):
        y_indices, x_indices = np.nonzero(labels_im == label)
        coords = np.column_stack((x_indices, y_indices))  
        all_coords.append(coords)
        
        center_x = int(np.mean(x_indices))
        center_y = int(np.mean(y_indices))
        centers.append(np.array([[center_x, center_y]]))  # 将中心点包装为 np.array([[x, y]])
    
    return centers

def merge_connected_components(binary_mask):
    kernel = np.ones((5, 5), np.uint8)  
    closed_mask = cv2.morphologyEx(binary_mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
    contours, _ = cv2.findContours(closed_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(closed_mask, contours, -1, (255), thickness=cv2.FILLED)
    return closed_mask

def describe_image(image, device, processor, model):
    inputs = processor(images=image, return_tensors="pt").to(device)
    output = model.generate(**inputs)
    description = processor.decode(output[0], skip_special_tokens=True)
    return description


def extract_concept(sentence, nlp):
    doc = nlp(sentence)
    subjects = []  # List to store subjects
    nouns = []     # List to store nouns
    for token in doc:
        # Extract subject(s)
        if "subj" in token.dep_:
            if token.text is not None:
                subjects.append(token.text)
        # Extract noun(s)
        if token.pos_ == "NOUN":
            nouns.append(token.text)
    # Return subjects if found, otherwise return nouns
    if subjects:
        return subjects
    elif nouns:
        return nouns
    else:
        return [sentence]  # Return the original sentence if no subject or noun is found


def has_preposition(scene):
    prepositions_pattern = r'\b(in|on|near|by|under|at|with|around|between)\b'
    return bool(re.search(prepositions_pattern, scene))

def add_preposition(scene):
    if has_preposition(scene) or scene == "":
        return scene
    default_prepositions = {
        "tree": "under",
        "beach": "near",
        "field": "in",
        "sun": "under",
        "sky": "in",
        "ball": "near",
        "mountain": "on",
        "room": "in"
    }
    
    for key, preposition in default_prepositions.items():
        if key in scene:
            scene = f"{preposition} {scene}"
            return scene
    
    scene = f"in {scene}"
    return scene

def generate_sentence(concept, attribute, action, scene):
    
    if len(attribute.split(" ")) == 1:
        templates = [
            "The {attribute} {concept} {action} {scene}.",  
            "{concept} {action} {scene} and looks {attribute}.",  
            "{concept} {action} {scene} while being {attribute}.",  
            "{scene}, the {attribute} {concept} {action}."  
        ]
    else:
        templates = [
            "The {concept}, {attribute}, {action} {scene}.", 
            "{concept} {attribute} is {action} {scene}.",     
            "{scene}, the {concept} {action} {attribute}.",    
            "{scene}, {concept} is {action} and {attribute}.",  
            "{concept} {action} {scene} with {attribute}."  
        ]
    
    template = random.choice(templates)
    sentence = template.format(attribute=attribute, concept=concept, action=action, scene=add_preposition(scene))
    return sentence

def clean_text(text_prompt):    
    allowed_punctuation = {',', '.', '?', '!'}
    cleaned_text = ''.join(c if c.isalnum() or c.isspace() or c in allowed_punctuation else ' ' for c in text_prompt)
    cleaned_text = re.sub(r'\s+', ' ', cleaned_text).strip()
    return cleaned_text