Danbooru Prompt Selector/TEST2024/NAIA_search.py

import pandas as pd
import re

# 모든 키워드를 포함하는 행만 필터링하는 함수
def filter_rows_containing_all_keywords(df, keywords):
    special_chars = r".^$*+?{}[]\|()"
    # 모든 키워드에 대한 boolean mask 초기화, df의 인덱스를 사용
    final_mask = pd.Series([True] * len(df), index=df.index)

    # 각 키워드에 대해 DataFrame의 모든 열을 검사하고 boolean mask 생성 및 저장
    for keyword in keywords:
        keyword_mask = pd.Series([False] * len(df), index=df.index)
        request_regex = False
        if any(char in keyword for char in special_chars): 
            keyword = re.escape(keyword)
            request_regex = True

        for column in ['copyright', 'character', 'artist', 'meta', 'general']:
            if df[column].dtype == 'object':
                if request_regex: keyword_mask |= df[column].str.contains(keyword, na=False, regex=True)
                else: keyword_mask |= df[column].str.contains(keyword, na=False)

        final_mask &= keyword_mask

    return df[final_mask]

def filter_rows_not_containing_all_keywords(df, keywords):
    special_chars = r".^$*+?{}[]\|()"
    # 모든 키워드를 포함하지 않는 행을 필터링하기 위한 boolean mask 초기화, df의 인덱스를 사용
    final_mask = pd.Series([True] * len(df), index=df.index)

    # 각 키워드에 대해 DataFrame의 모든 열을 검사하고 boolean mask 생성 및 저장
    for keyword in keywords:
        keyword_mask = pd.Series([False] * len(df), index=df.index)
        request_regex = False
        if any(char in keyword for char in special_chars): 
            keyword = re.escape(keyword)
            request_regex = True
        for column in ['copyright', 'character', 'artist', 'meta', 'general']:
            if df[column].dtype == 'object':
                if request_regex: keyword_mask |= df[column].str.contains(keyword, na=False, regex=True)
                else: keyword_mask |= df[column].str.contains(keyword, na=False)

        # 모든 키워드를 포함하는 행에 대한 mask를 반전시켜 final_mask에 저장
        final_mask &= ~keyword_mask

    return df[final_mask]

def process_asterisk_group(df, asterisk_group):
    special_chars = r".^$*+?{}[]\|()"
    # 각 키워드 앞의 '*'를 제거하고 맨 뒤에 ',' 추가
    asterisk_keywords = [keyword.lstrip('*') + ',' for keyword in asterisk_group]

    # 각 행에 대해 임시 문자열 search_string을 만들고 검색 수행
    df['search_string'] = df[['copyright', 'character', 'artist', 'meta', 'general']].apply(lambda x: ' ' + ', '.join(x.astype(str)) + ',', axis=1)
    for keyword in asterisk_keywords:
        request_regex = False
        if any(char in keyword for char in special_chars): 
            keyword = re.escape(keyword)
            request_regex = True
        if request_regex: df = df[df['search_string'].str.contains(keyword, na=False, regex=True)]
        else:  df = df[df['search_string'].str.contains(keyword, na=False)]
    df.drop('search_string', axis=1, inplace=True)

    return df

def process_perfect_negative_group(df, perfect_negative_group):
    special_chars = r".^$*+?{}[]\|()"
    # 각 키워드 앞의 '~'를 제거하고 맨 뒤에 ',' 추가
    perfect_negative_keywords = [keyword.lstrip('~') + ',' for keyword in perfect_negative_group]

    # 각 행에 대해 임시 문자열 search_string을 만듦
    df['search_string'] = df[['copyright', 'character', 'artist', 'meta', 'general']].apply(lambda x: ' ' + ', '.join(x.astype(str)) + ',', axis=1)

    # 모든 키워드에 대한 검색 결과를 하나의 boolean Series로 결합
    combined_mask = pd.Series([True] * len(df), index=df.index)
    for keyword in perfect_negative_keywords:
        request_regex = False
        if any(char in keyword for char in special_chars): 
            keyword = re.escape(keyword)
            request_regex = True
        if request_regex: keyword_mask = df['search_string'].str.contains(keyword, na=False, regex=True)
        else: df['search_string'].str.contains(keyword, na=False)
        combined_mask &= ~keyword_mask

    # 최종적으로 일치하지 않는 행만 필터링
    df = df[combined_mask]

    # search_string 열 제거
    df.drop('search_string', axis=1, inplace=True)

    return df

def search(df, search_request, exclude_request, E, N, S, G):
    if(E == 0):
        df = df[~(df['rating'] == 'e')]
    if(N == 0):
        df = df[~(df['rating'] == 'q')]
    if(S == 0):
        df = df[~(df['rating'] == 's')]
    if(G == 0):
        df = df[~(df['rating'] == 'g')]
    if(len(df) == 0):
        return None
    
    special_chars = r".^$*+?{}[]\|()"

    #search_request에 대한 처리
    #처리순서 normal -> curly -> asterisk
    split_requests = [item.strip() for item in search_request.split(',')]

    curly_brace_group = [item for item in split_requests if item.startswith('{') and item.endswith('}')]
    asterisk_group = [item for item in split_requests if item.startswith('*')]
    normal_group = [item for item in split_requests if item not in curly_brace_group + asterisk_group]
    #normal_group = [re.escape(item) if any(char in item for char in special_chars) else item for item in normal_group]
    negative_split_requests = [item.strip() for item in exclude_request.split(',')]
    perfect_negative_group = [item for item in negative_split_requests if item.startswith('~')]
    negative_group = [item for item in negative_split_requests if item not in perfect_negative_group]
    #negative_group = [re.escape(item) if any(char in item for char in special_chars) else item for item in negative_group]

    if '' in split_requests:
        split_requests.remove('')
    if '' in negative_split_requests:
        negative_split_requests.remove('')

    #포지티브
    if split_requests:
        #normal 처리
        if normal_group:
            df = filter_rows_containing_all_keywords(df, normal_group)
            if(len(df) == 0):
                return None

        #OR 처리
        if curly_brace_group:
            for keyword in curly_brace_group:
                or_search_keyword = [item.strip() for item in keyword[1:-1].split('|')]
                results = pd.DataFrame()
                for keyword in or_search_keyword:
                    if keyword.startswith('*'):
                        keyword = keyword[1:]
                    request_regex = False
                    if any(char in keyword for char in special_chars): 
                        keyword = re.escape(keyword)
                        request_regex = True
                    for column in ['copyright', 'character', 'artist', 'meta', 'general']:
                        if request_regex: matched_rows = df[df[column].str.contains(keyword, na=False, regex=True)]
                        else: matched_rows = df[df[column].str.contains(keyword, na=False)]
                        if not matched_rows.empty:
                            results = pd.concat([results, matched_rows])
                            break
                del[[df]]
                df = results.copy()
                del[[results]]
            if(len(df) == 0):
                return None
        
        #Perfect Matching 처리
        if asterisk_group:
            df = process_asterisk_group(df,asterisk_group)
            if(len(df) == 0):
                return None
    
    #Exclude 처리
    if negative_split_requests:
        if negative_group:
            df = filter_rows_not_containing_all_keywords(df, negative_group)
            if(len(df) == 0):
                return None
        
        if perfect_negative_group:
            df = process_perfect_negative_group(df, perfect_negative_group)
            if(len(df) == 0):
                return None
    return df