FindCenter.py

import cv2
import numpy as np
import os
import glob

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
INPUT_FOLDER = os.path.join(BASE_DIR, 'image')
OUTPUT_FOLDER = os.path.join(BASE_DIR, 'output')

CROP_WIDTH = 96
CROP_HEIGHT = 96

def find_character_center(image_path, output_dir):
    img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    if img is None: return None
    
    h, w, c = img.shape
    if c < 4: return None
        
    alpha = img[:, :, 3]
    _, binary = cv2.threshold(alpha, 254, 1, cv2.THRESH_BINARY)
    if np.max(binary) == 0: return None

    win_w, win_h = 50, 70
    head_h = 30
    body_w = 20
    
    kernel = np.zeros((win_h, win_w), dtype=np.float32)
    kernel[0:head_h, :] = 1.0
    body_start_x = (win_w - body_w) // 2
    kernel[head_h:win_h, body_start_x:body_start_x + body_w] = 2.5

    density_map = cv2.filter2D(binary.astype(float), -1, kernel)
    
    y_indices, x_indices = np.indices((h, w))
    center_y_img, center_x_img = h / 2, w / 2
    dist_from_center = np.sqrt((x_indices - center_x_img)**2 + (y_indices - center_y_img)**2)
    max_dist = np.sqrt(center_x_img**2 + center_y_img**2)
    center_weight = 1.0 - 0.3 * (dist_from_center / max_dist) 
    bottom_weight = 1.0 + 0.15 * (y_indices / h)
    
    weight_map = center_weight * bottom_weight
    weighted_density = density_map * weight_map
    
    max_val = np.max(weighted_density)
    best_y, best_x = np.where(weighted_density >= max_val * 0.999)
    
    if len(best_x) == 0:
        return None
        
    rough_x = int(np.mean(best_x))

    spine_half_width = body_w // 2
    col_start = max(0, rough_x - spine_half_width)
    col_end = min(w, rough_x + spine_half_width)
    
    body_column = binary[:, col_start:col_end]
    y_coords_in_spine, x_coords_in_spine = np.where(body_column == 1)
    
    if len(y_coords_in_spine) > 0:
        center_x = col_start + int(np.median(x_coords_in_spine))
        
        y_proj = np.sum(body_column, axis=1)
        rough_y = int(np.mean(best_y))
        
        foot_y = rough_y
        while foot_y < h and y_proj[foot_y] > 0:
            foot_y += 1
            
        true_foot_y = foot_y - 1
        center_y = true_foot_y - (win_h // 2)
    else:
        center_x = rough_x
        center_y = int(np.mean(best_y))

    start_x = center_x - CROP_WIDTH // 2
    end_x = start_x + CROP_WIDTH
    start_y = center_y - CROP_HEIGHT // 2
    end_y = start_y + CROP_HEIGHT
    
    src_start_x = max(0, start_x)
    src_start_y = max(0, start_y)
    src_end_x = min(w, end_x)
    src_end_y = min(h, end_y)
    
    dst_start_x = max(0, -start_x)
    dst_start_y = max(0, -start_y)
    dst_end_x = dst_start_x + (src_end_x - src_start_x)
    dst_end_y = dst_start_y + (src_end_y - src_start_y)
    
    cropped_img = np.zeros((CROP_HEIGHT, CROP_WIDTH, 4), dtype=np.uint8)
    if src_end_x > src_start_x and src_end_y > src_start_y:
        cropped_img[dst_start_y:dst_end_y, dst_start_x:dst_end_x] = img[src_start_y:src_end_y, src_start_x:src_end_x]

    filename = os.path.basename(image_path)
    output_path = os.path.join(output_dir, filename)
    cv2.imwrite(output_path, cropped_img)
    
    return (center_x, center_y)

def process_directory(input_dir, output_dir):
    if not os.path.exists(input_dir):
        print(f"'{input_dir}' 폴더가 존재하지 않습니다.")
        return False
        
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    image_files = glob.glob(os.path.join(input_dir, '*.png'))
    if not image_files:
        print(f"'{input_dir}' 폴더에 PNG 사진 파일이 없습니다.")
        return False

    for file_path in image_files:
        filename = os.path.basename(file_path)
        center = find_character_center(file_path, output_dir)
        if center:
            print(f"{filename}")
            
    return True

if __name__ == "__main__":
    try:
        process_directory(INPUT_FOLDER, OUTPUT_FOLDER)
    except Exception as e:
        print(f"\n{e}")
    finally: 
        input("\n끗")