import gradio as gr
import cv2
import numpy as np
import matplotlib.pyplot as plt
import time
import os
from sklearn.cluster import MiniBatchKMeans
from PIL import Image
from transformers import ViTImageProcessor, ViTForImageClassification
import torch
from numba import jit, prange

# 设置环境变量以防止内存泄漏
os.environ["OMP_NUM_THREADS"] = "1"

# 加载预训练的 ViT 模型和图像处理器
model_name = "google/vit-base-patch16-224"
image_processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

@jit(nopython=True, parallel=True)
def rgb_to_gray_numba(image):
    """使用 Numba JIT 编译器将 RGB 图像转换为灰度图像"""
    h, w, _ = image.shape
    gray_image = np.empty((h, w), dtype=np.uint8)
    for i in prange(h):
        for j in prange(w):
            r, g, b = image[i, j]
            gray_image[i, j] = 0.2989 * r + 0.5870 * g + 0.1140 * b
    return gray_image

def create_color_palette(colors, size=100):
    """创建颜色调色板"""
    palette = np.zeros((4 * size, 4 * size, 3), dtype=np.uint8)
    for i in range(4):
        for j in range(4):
            palette[i * size:(i + 1) * size, j * size:(j + 1) * size] = colors[i * 4 + j]
    return palette

def identify_image(image):
    """识别图像中的对象"""
    inputs = image_processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
    top_5 = torch.topk(probabilities, 5)

    results = []
    for i in range(5):
        label = model.config.id2label[top_5.indices[i].item()]
        score = top_5.values[i].item()
        results.append(f"{label}: {score:.2f}")
    
    return ", ".join(results)

def analyze_image(image):
    overall_start_time = time.time()
    times = {}

    # 0. 图像分类
    start_time = time.time()
    top_5_classes = identify_image(image)
    times['classification'] = f"{time.time() - start_time:.3f} seconds"

    # 将图像转换为 NumPy 数组以进行 OpenCV 操作
    image_np = np.array(image)

    # 1. 直方图
    start_time = time.time()
    gray_image = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    hist = cv2.calcHist([gray_image], [0], None, [256], [0, 256])
    hist = cv2.normalize(hist, hist).flatten()
    
    fig, ax = plt.subplots(figsize=(2.56, 1), dpi=100)
    ax.plot(hist, color='black')
    ax.set_xlim([0, 256])
    ax.set_ylim([0, 1])
    ax.axis('off')
    
    histogram_path = "histogram.png"
    plt.savefig(histogram_path, format='png', bbox_inches='tight', pad_inches=0)
    plt.close(fig)
    times['histogram'] = f"{time.time() - start_time:.3f} seconds"

    # 2. 边缘检测
    start_time = time.time()
    edges = cv2.Canny(gray_image, 100, 200)
    edges_path = "edges.png"
    cv2.imwrite(edges_path, edges)
    times['edge_detection'] = f"{time.time() - start_time:.3f} seconds"

    # 3. 主颜色提取（16种颜色）
    start_time = time.time()
    
    # 调整图像大小以加快处理速度
    small_image = cv2.resize(image_np, (100, 100))
    pixels = small_image.reshape((-1, 3))
    pixels = np.float32(pixels)

    # 使用 MiniBatchKMeans 代替 K-means
    kmeans = MiniBatchKMeans(n_clusters=16, batch_size=100, max_iter=10)
    kmeans.fit(pixels)
    colors = kmeans.cluster_centers_
    
    # 创建颜色调色板
    color_palette = create_color_palette(colors)
    palette_path = "palette.png"
    cv2.imwrite(palette_path, color_palette)
    times['main_colors'] = f"{time.time() - start_time:.3f} seconds"

    # 计算总体耗时
    overall_elapsed_time = time.time() - overall_start_time
    times['overall'] = f"{overall_elapsed_time:.3f} seconds"

    # 准备时间输出
    timing_output = "\n".join([f"{k}: {v}" for k, v in times.items()])

    return top_5_classes, histogram_path, edges_path, palette_path, timing_output

# 定义 Gradio 界面
interface = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil", label="上传图像"),
    outputs=[
        gr.Textbox(label="图像分类 (Top 5)"),
        gr.Image(label="灰度直方图"),
        gr.Image(label="边缘检测"),
        gr.Image(label="颜色调色板"),
        gr.Textbox(label="处理时间")
    ],
    title="高级图像分析",
    description="上传图像以分析其分类、直方图、边缘和主要颜色。",
    allow_flagging="never"
)

# 启动界面
interface.launch(share=True)