app.py

import os  
from flask import Flask, request, jsonify
from PIL import Image, ImageDraw, ImageFont
import numpy as np
from ultralytics import YOLO
from collections import Counter
import boto3
from io import BytesIO
from dotenv import load_dotenv

# .env 파일 로드
load_dotenv()

app = Flask(__name__)

model = YOLO('yolov8n.pt')

# 환경 변수에서 AWS S3 설정 가져오기
AWS_ACCESS_KEY_ID = os.getenv('AWS_ACCESS_KEY_ID')
AWS_SECRET_ACCESS_KEY = os.getenv('AWS_SECRET_ACCESS_KEY')
AWS_REGION = os.getenv('AWS_REGION')
S3_BUCKET_NAME = os.getenv('S3_BUCKET_NAME')

# AWS S3 설정
s3_client = boto3.client(
    's3',
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
    region_name=AWS_REGION
)

S3_BUCKET_NAME = 'boda-bucket'  

@app.route('/')
def hello_world():
    return 'Hello World!'

@app.route('/detect', methods=['POST'])
def detect_objects():
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400

    file = request.files['file']

    try:
        # 업로드된 파일을 이미지로 변환
        input_image = Image.open(file).convert("RGB")
        print(f"Original image size: {input_image.size}")

        # 가로와 세로를 변경 (90도 회전)
        input_image = input_image.transpose(Image.ROTATE_270)
        print(f"Rotated image size: {input_image.size}")
    except Exception as e:
        return jsonify({'error': f'Invalid image format: {str(e)}'}), 400

    # YOLOv8 모델로 객체 탐지
    results = model(input_image)
    detections = results[0].boxes.data.cpu().numpy() 

    # 신뢰도
    confidence_threshold = 0.5
    
    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
    font_path_korean = "/usr/share/fonts/truetype/nanum/NanumGothic.ttf"

    try:
        font = ImageFont.truetype(font_path_korean, 40) 
    except IOError:
        print(f"Font not found at {font_path_korean}. Using default font.")
        font = ImageFont.load_default()

    
    # 결과를 JSON 형식으로 변환
    detected_objects = []
    for detection in detections:
        # numpy 배열의 값을 직접 사용
        x_min, y_min, x_max, y_max = detection[:4]
        confidence = detection[4] if len(detection) > 4 else None
        class_idx = detection[5] if len(detection) > 5 else None
        
        class_name = model.names.get(int(class_idx), "Unknown")
        
        print(f"Detected object: {x_min}, {y_min}, {x_max}, {y_max}, confidence: {confidence}, class: {class_name}")
        
        if confidence is not None and confidence < confidence_threshold:
            continue
        
        class_name = model.names.get(int(class_idx), "Unknown")
        
        # 박스 내부 이미지 추출
        cropped_box = input_image.crop((x_min, y_min, x_max, y_max))
        # 주요 색상 추출
        main_color_rgb = get_main_color(cropped_box)
        main_color_name = rgb_to_color_name(main_color_rgb) 

        print(f"Detected object: {x_min}, {y_min}, {x_max}, {y_max}, confidence: {confidence}, class: {class_name}, main color: {main_color_name}")

        detected_object = {
            'bounding_box': {
                'x_min': int(x_min),
                'y_min': int(y_min),
                'x_max': int(x_max),
                'y_max': int(y_max)
            }
        }
        if confidence is not None:
            detected_object['confidence'] = round(float(confidence), 2)
        if class_idx is not None:
            detected_object['class_id'] = int(class_idx)
            detected_object['class_name'] = class_name


        detected_objects.append(detected_object)
        # 이미지 위에 박스 그리기
        draw = ImageDraw.Draw(input_image)
        draw.rectangle([x_min, y_min, x_max, y_max], outline=main_color_rgb, width=12)

        # 클래스 이름과 신뢰도를 박스 위에 추가 (옵션)
        if confidence is not None and class_idx is not None:
            label = f"{main_color_name}"
                # 텍스트 경계 상자 계산
            bbox = draw.textbbox((0, 0), label, font=font)  # (0, 0)은 임시 위치
            text_width = bbox[2] - bbox[0]  # 너비 계산
            text_height = bbox[3] - bbox[1]  # 높이 계산

            # 텍스트의 최적 위치 계산 (텍스트가 박스 위에 표시되도록 조정)
            text_x = x_min  # 텍스트는 박스의 왼쪽 상단 x_min에서 시작
            text_y = max(0, y_min - text_height - 5)  # 텍스트가 이미지 경계를 벗어나지 않도록 보정

            text_position = (text_x, text_y)

            draw.rectangle(
                [text_x, text_y, text_x + text_width, text_y + text_height],
                fill='white',
            )

            draw.text(text_position, label, fill="black", font=font)

    try:
        s3_url = save_image_to_s3(input_image, 'detected_image.png')
        print(f"Image uploaded to S3: {s3_url}")
    except Exception as e:
        return jsonify({'error': f'Failed to upload image to S3: {str(e)}'}), 500


    # 결과 이미지를 저장
    output_image_path = "./detected_image1.png"  # 저장할 이미지 경로
    input_image.save(output_image_path)
    print(f"Image saved with detections: {output_image_path}")

    response = {
        'status': 'success',
        'image':s3_url
        # 'total_detections': len(detected_objects)
    }
    print(response);

    return jsonify(response), 200

def save_image_to_s3(image, filename):
    """
    이미지 파일을 S3에 업로드하는 함수
    """
    buffer = BytesIO()
    image.save(buffer, format='PNG')  # 이미지를 버퍼에 저장
    buffer.seek(0)

    s3_client.upload_fileobj(
        buffer,
        S3_BUCKET_NAME,
        filename,
        ExtraArgs={'ContentType': 'image/png'}
    )

    s3_url = f"https://{S3_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{filename}"
    return s3_url


def get_main_color(image, num_colors=1):
    """Extract the main color from an image."""
    
    if image.mode != "RGB":
        image = image.convert("RGB")
        
        # 이미지 크기 가져오기
    width, height = image.size

    # 중앙 부분 잘라내기 (가로 반, 세로 반)
    left = width // 4
    upper = height // 4
    right = 3 * (width // 4)
    lower = 3 * (height // 4)
    cropped_image = image.crop((left, upper, right, lower))

    # 크롭된 이미지를 50x50으로 축소
    cropped_image = cropped_image.resize((50, 50))
        
    # image = image.resize((50, 50))  
    pixels = np.array(cropped_image).reshape(-1, 3)  
    counter = Counter([tuple(pixel) for pixel in pixels])
    most_common_color = counter.most_common(1)[0][0]  
    return most_common_color


def rgb_to_color_name(rgb):
    """Convert RGB to a refined color name with range conditions."""
    r, g, b = rgb
    print(f"RGB: {r}, {g}, {b}")

    if 0 <= r <= 84 and 0 <= g <= 84 and 0 <= b <= 84:
        return "검정색"
    elif 84 <= r <= 171 and 0 <= g <= 84 and 0 <= b <= 84:
        return "갈색"
    elif 171 <= r <= 255 and 0 <= g <= 84 and 0 <= b <= 84:
        return "빨간색"
    elif 0 <= r <= 84 and 84 <= g <= 171 and 0 <= b <= 84:
        return "진한 초록색"
    elif 0 <= r <= 84 and 171 <= g <= 255 and 0 <= b <= 84:
        return "초록색"
    elif 84 <= r <= 171 and 84 <= g <= 171 and 0 <= b <= 84:
        return "카키색"
    elif 171 <= r <= 255 and 84 <= g <= 171 and 0 <= b <= 84:
        return "주황색"
    elif 85 <= r <= 171 and 171 <= g <= 255 and 0 <= b <= 84:
        return "연두색"
    elif 171 <= r <= 255 and 171 <= g <= 255 and 0 <= b <= 84:
        return "노란색"
    
    
    elif 0 <= r <= 84 and 0 <= g <= 84 and 84 <= b <= 171:
        return "남색"
    elif 0 <= r <= 84 and 84 <= g <= 171 and 84 <= b <= 171:
        return "짙은 하늘색"
    elif 0 <= r <= 84 and 171 <= g <= 255 and 84 <= b <= 171:
        return "형광 초록색"
    elif 84 <= r <= 171 and 0 <= g <= 84 and 84 <= b <= 171:
        return "자주색"
    elif 84 <= r <= 171 and 84 <= g <= 171 and 84 <= b <= 171:
        return "회색"
    elif 84 <= r <= 171 and 171 <= g <= 255 and 84 <= b <= 171:
        return "짙은 연두색"
    elif 171 <= r <= 255 and 0 <= g <= 84 and 84 <= b <= 171:
        return "핑크색"
    elif 171 <= r <= 255 and 84 <= g <= 171 and 84 <= b <= 171:
        return "연한 갈색"
    elif 171 <= r <= 255 and 171 <= g <= 255 and 84 <= b <= 171:
        return "연한 노란색"
    
    elif 0 <= r <= 84 and 0 <= g <= 84 and 171 <= b <= 255:
        return "파랑색"
    elif 0 <= r <= 84 and 84 <= g <= 171 and 171 <= b <= 255:
        return "하늘색"
    elif 0 <= r <= 84 and 171 <= g <= 255 and 171 <= b <= 255:
        return "청록색"
    elif 84 <= r <= 171 and 0 <= g <= 84 and 171 <= b <= 255:
        return "보라색"
    elif 84 <= r <= 171 and 84 <= g <= 171 and 171 <= b <= 255:
        return "연보라색"
    elif 84 <= r <= 171 and 171 <= g <= 255 and 171 <= b <= 255:
        return "소라색"
    elif 171 <= r <= 255 and 0 <= g <= 84 and 171 <= b <= 255:
        return "형광 보라색"
    elif 171 <= r <= 255 and 84 <= g <= 171 and 171 <= b <= 255:
        return "연핑크색"
    elif 171 <= r <= 255 and 171 <= g <= 255 and 171 <= b <= 255:
        return "흰색"

    else:
        return "기타 색상"

# def rgb_to_color_name(rgb):
#     r, g, b = rgb
#     closest_color = None
#     min_distance = float('inf')

#     for color_name, (cr, cg, cb) in colors.items():
#         # Calculate Euclidean distance
#         distance = ((r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2) ** 0.5
#         if distance < min_distance:
#             min_distance = distance
#             closest_color = color_name

#     return closest_color

# Define the colors dictionary with central RGB values for each range
colors = {
    "검정색": (42, 42, 42),           # Central value for 0-84 range
    "갈색": (128, 42, 42),           # Central value for 84-171 range
    "빨간색": (213, 42, 42),         # Central value for 171-255 range
    "진한 초록색": (42, 128, 42),
    "초록색": (42, 213, 42),
    "카키색": (128, 128, 42),
    "주황색": (213, 128, 42),
    "연두색": (128, 213, 42),
    "노란색": (213, 213, 42),
    "남색": (42, 42, 128),
    "짙은 하늘색": (42, 128, 128),
    "형광 초록색": (42, 213, 128),
    "자주색": (128, 42, 128),
    "회색": (128, 128, 128),
    "짙은 연두색": (128, 213, 128),
    "핑크색": (213, 42, 128),
    "연한 갈색": (213, 128, 128),
    "연한 노란색": (213, 213, 128),
    "파랑색": (42, 42, 213),
    "하늘색": (42, 128, 213),
    "청록색": (42, 213, 213),
    "보라색": (128, 42, 213),
    "연보라색": (128, 128, 213),
    "소라색": (128, 213, 213),
    "형광 보라색": (213, 42, 213),
    "연핑크색": (213, 128, 213),
    "흰색": (213, 213, 213)
}

if __name__ == '__main__':
    app.run('0.0.0.0', port=8080, debug=True)