R360-UAVmodelTool/image_classifier.py

import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from PIL import Image, ImageTk, ImageEnhance
import os
import shutil
from pathlib import Path
import numpy as np
from collections import Counter
import cv2

class ImageClassifier:
    def __init__(self, root):
        self.root = root
        self.root.title("图片分类器")
        self.root.geometry("1200x800")
        
        # 当前图片路径
        self.current_image_path = None
        # 当前文件夹路径
        self.current_folder_path = None
        # 图片列表
        self.image_list = []
        # 当前图片索引
        self.current_index = 0
        # 图片处理参数
        self.zoom_factor = 1.0
        self.rotation_angle = 0
        self.brightness_factor = 1.0
        # 天空检测参数
        self.sky_threshold = 0.7  # 天空占比阈值
        self.blue_threshold = 0.6  # 蓝色占比阈值
        # 操作历史记录
        self.operation_history = []  # 存储操作历史
        
        self.setup_ui()
        self.setup_keyboard_shortcuts()
        
    def setup_ui(self):
        # 创建主框架
        main_frame = ttk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 左侧控制面板
        control_frame = ttk.Frame(main_frame)
        control_frame.pack(side=tk.LEFT, fill=tk.Y, padx=5)
        
        # 选择文件夹按钮
        ttk.Button(control_frame, text="选择文件夹 (F)", command=self.select_folder).pack(pady=5)
        
        # 分类按钮
        ttk.Button(control_frame, text="有无人机 (1)", command=lambda: self.classify_image("drone")).pack(pady=5)
        ttk.Button(control_frame, text="有鸟类 (2)", command=lambda: self.classify_image("bird")).pack(pady=5)
        ttk.Button(control_frame, text="有鸟类和无人机 (3)", command=lambda: self.classify_image("both")).pack(pady=5)
        ttk.Button(control_frame, text="无目标 (4)", command=lambda: self.classify_image("none")).pack(pady=5)
        
        # 导航按钮
        nav_frame = ttk.Frame(control_frame)
        nav_frame.pack(pady=10)
        ttk.Button(nav_frame, text="上一张 (←)", command=self.prev_image).pack(side=tk.LEFT, padx=5)
        ttk.Button(nav_frame, text="下一张 (→)", command=self.next_image).pack(side=tk.LEFT, padx=5)
        
        # 撤回按钮
        ttk.Button(control_frame, text="撤回操作 (Z)", command=self.undo_operation).pack(pady=5)
        
        # 图片处理控制
        process_frame = ttk.LabelFrame(control_frame, text="图片处理")
        process_frame.pack(pady=10, fill=tk.X)
        
        # 缩放控制
        ttk.Label(process_frame, text="缩放:").pack()
        self.zoom_scale = ttk.Scale(process_frame, from_=0.1, to=3.0, orient=tk.HORIZONTAL, 
                                  command=self.update_zoom)
        self.zoom_scale.set(1.0)
        self.zoom_scale.pack(fill=tk.X, padx=5)
        
        # 旋转控制
        ttk.Label(process_frame, text="旋转:").pack()
        self.rotation_scale = ttk.Scale(process_frame, from_=0, to=360, orient=tk.HORIZONTAL,
                                      command=self.update_rotation)
        self.rotation_scale.set(0)
        self.rotation_scale.pack(fill=tk.X, padx=5)
        
        # 亮度控制
        ttk.Label(process_frame, text="亮度:").pack()
        self.brightness_scale = ttk.Scale(process_frame, from_=0.1, to=2.0, orient=tk.HORIZONTAL,
                                        command=self.update_brightness)
        self.brightness_scale.set(1.0)
        self.brightness_scale.pack(fill=tk.X, padx=5)
        
        # 天空检测
        ttk.Label(process_frame, text="天空检测:").pack()
        self.auto_detect_var = tk.BooleanVar(value=False)
        ttk.Checkbutton(process_frame, text="自动检测天空图片", 
                       variable=self.auto_detect_var,
                       command=self.toggle_auto_detect).pack()
        
        # 天空阈值控制
        ttk.Label(process_frame, text="天空阈值:").pack()
        self.threshold_scale = ttk.Scale(process_frame, from_=0.5, to=0.95, orient=tk.HORIZONTAL,
                                       command=self.update_threshold)
        self.threshold_scale.set(0.7)
        self.threshold_scale.pack(fill=tk.X, padx=5)
        
        # 重置按钮
        ttk.Button(process_frame, text="重置图片 (R)", command=self.reset_image).pack(pady=5)
        
        # 右侧图片显示区域
        self.image_frame = ttk.Frame(main_frame)
        self.image_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        
        # 图片标签
        self.image_label = ttk.Label(self.image_frame)
        self.image_label.pack(fill=tk.BOTH, expand=True)
        
        # 状态标签
        self.status_label = ttk.Label(self.root, text="请选择文件夹")
        self.status_label.pack(side=tk.BOTTOM, pady=5)
        
        # 添加图片信息显示
        self.info_label = ttk.Label(self.root, text="")
        self.info_label.pack(side=tk.BOTTOM, pady=5)
        
    def setup_keyboard_shortcuts(self):
        self.root.bind('<F5>', lambda e: self.select_folder())
        self.root.bind('<Left>', lambda e: self.prev_image())
        self.root.bind('<Right>', lambda e: self.next_image())
        self.root.bind('1', lambda e: self.classify_image("drone"))
        self.root.bind('2', lambda e: self.classify_image("bird"))
        self.root.bind('3', lambda e: self.classify_image("both"))
        self.root.bind('4', lambda e: self.classify_image("none"))
        self.root.bind('r', lambda e: self.reset_image())
        self.root.bind('z', lambda e: self.undo_operation())
        
    def update_zoom(self, value):
        self.zoom_factor = float(value)
        self.show_current_image()
        
    def update_rotation(self, value):
        self.rotation_angle = float(value)
        self.show_current_image()
        
    def update_brightness(self, value):
        self.brightness_factor = float(value)
        self.show_current_image()
        
    def reset_image(self):
        self.zoom_factor = 1.0
        self.rotation_angle = 0
        self.brightness_factor = 1.0
        self.zoom_scale.set(1.0)
        self.rotation_scale.set(0)
        self.brightness_scale.set(1.0)
        self.show_current_image()
        
    def select_folder(self):
        folder_path = filedialog.askdirectory()
        if folder_path:
            self.current_folder_path = folder_path
            self.load_images()
            
    def load_images(self):
        self.image_list = []
        for file in os.listdir(self.current_folder_path):
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                self.image_list.append(os.path.join(self.current_folder_path, file))
        
        if self.image_list:
            self.current_index = 0
            self.show_current_image()
            self.status_label.config(text=f"已加载 {len(self.image_list)} 张图片")
        else:
            messagebox.showinfo("提示", "所选文件夹中没有图片")
            
    def show_current_image(self):
        if 0 <= self.current_index < len(self.image_list):
            self.current_image_path = self.image_list[self.current_index]
            
            # 打开图片
            image = Image.open(self.current_image_path)
            
            # 如果启用了自动检测，检查是否为天空图片
            if self.auto_detect_var.get():
                if self.is_sky_image(image):
                    self.classify_image("none")
                    return
            
            # 应用亮度调整
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(self.brightness_factor)
            
            # 应用旋转
            image = image.rotate(self.rotation_angle, expand=True)
            
            # 计算调整后的尺寸，保持宽高比
            width, height = image.size
            max_size = (800, 600)
            ratio = min(max_size[0]/width, max_size[1]/height)
            new_size = (int(width*ratio*self.zoom_factor), int(height*ratio*self.zoom_factor))
            image = image.resize(new_size, Image.Resampling.LANCZOS)
            
            # 转换为PhotoImage
            photo = ImageTk.PhotoImage(image)
            
            # 更新图片显示
            self.image_label.config(image=photo)
            self.image_label.image = photo  # 保持引用
            
            # 更新状态
            self.status_label.config(text=f"图片 {self.current_index + 1}/{len(self.image_list)}")
            
            # 显示图片信息
            file_size = os.path.getsize(self.current_image_path) / 1024  # 转换为KB
            self.info_label.config(text=f"文件名: {os.path.basename(self.current_image_path)} | 尺寸: {width}x{height} | 大小: {file_size:.1f}KB")
            
    def next_image(self):
        if self.image_list:
            self.current_index = (self.current_index + 1) % len(self.image_list)
            self.show_current_image()
            
    def prev_image(self):
        if self.image_list:
            self.current_index = (self.current_index - 1) % len(self.image_list)
            self.show_current_image()
            
    def classify_image(self, category):
        if not self.current_image_path:
            return
            
        # 创建分类文件夹
        category_folders = {
            "drone": "有无人机",
            "bird": "有鸟类",
            "both": "有鸟类和无人机",
            "none": "无目标"
        }
        
        target_folder = os.path.join(self.current_folder_path, category_folders[category])
        os.makedirs(target_folder, exist_ok=True)
        
        # 移动图片到对应文件夹
        filename = os.path.basename(self.current_image_path)
        target_path = os.path.join(target_folder, filename)
        
        try:
            # 记录操作历史（在移动之前）
            operation_record = {
                'action': 'classify',
                'original_path': self.current_image_path,
                'target_path': target_path,
                'category': category,
                'index': self.current_index,
                'filename': filename
            }
            
            shutil.move(self.current_image_path, target_path)
            
            # 添加到操作历史
            self.operation_history.append(operation_record)
            
            # 限制历史记录数量（最多保留50条）
            if len(self.operation_history) > 50:
                self.operation_history.pop(0)
            
            self.image_list.pop(self.current_index)
            if self.image_list:
                self.current_index = self.current_index % len(self.image_list)
                self.show_current_image()
            else:
                self.image_label.config(image='')
                self.status_label.config(text="所有图片已分类完成")
                self.info_label.config(text="")
        except Exception as e:
            messagebox.showerror("错误", f"移动文件时出错：{str(e)}")

    def is_sky_image(self, image):
        # 转换为OpenCV格式
        img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        
        # 转换为HSV颜色空间
        hsv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2HSV)
        
        # 定义天空的HSV范围（包括蓝色和灰色）
        # 蓝色天空范围
        lower_blue = np.array([100, 50, 50])
        upper_blue = np.array([130, 255, 255])
        
        # 灰色天空范围
        lower_gray = np.array([0, 0, 50])
        upper_gray = np.array([180, 30, 200])
        
        # 创建掩码
        mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)
        mask_gray = cv2.inRange(hsv, lower_gray, upper_gray)
        
        # 合并掩码
        mask = cv2.bitwise_or(mask_blue, mask_gray)
        
        # 计算天空像素占比
        sky_ratio = np.sum(mask > 0) / (mask.shape[0] * mask.shape[1])
        
        # 计算图片上半部分的天空占比
        upper_half = mask[:mask.shape[0]//2, :]
        upper_sky_ratio = np.sum(upper_half > 0) / (upper_half.shape[0] * upper_half.shape[1])
        
        # 计算图片的纹理特征
        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        texture = cv2.Laplacian(gray, cv2.CV_64F).var()
        
        # 判断是否为天空图片
        # 1. 整体天空占比超过阈值
        # 2. 上半部分主要是天空
        # 3. 纹理变化较小（天空通常比较均匀）
        return (sky_ratio > self.sky_threshold and 
                upper_sky_ratio > self.blue_threshold and 
                texture < 500)  # 纹理阈值可以根据需要调整
        
    def toggle_auto_detect(self):
        if self.auto_detect_var.get() and self.current_image_path:
            if self.is_sky_image(Image.open(self.current_image_path)):
                self.classify_image("none")
                messagebox.showinfo("提示", "检测到天空图片，已自动分类为'无目标'")
                
    def update_threshold(self, value):
        self.sky_threshold = float(value)
        
    def undo_operation(self):
        """撤回最后一次分类操作"""
        if not self.operation_history:
            messagebox.showinfo("提示", "没有可撤回的操作")
            return
            
        # 获取最后一次操作
        last_operation = self.operation_history.pop()
        
        if last_operation['action'] == 'classify':
            try:
                # 将文件移回原位置
                shutil.move(last_operation['target_path'], last_operation['original_path'])
                
                # 重新加载图片列表
                self.load_images()
                
                # 尝试定位到撤回的图片
                try:
                    restored_index = self.image_list.index(last_operation['original_path'])
                    self.current_index = restored_index
                except ValueError:
                    # 如果找不到，就显示第一张
                    self.current_index = 0
                
                if self.image_list:
                    self.show_current_image()
                    messagebox.showinfo("成功", f"已撤回对 '{last_operation['filename']}' 的分类操作")
                else:
                    messagebox.showinfo("提示", "撤回成功，但当前文件夹中没有图片")
                    
            except Exception as e:
                messagebox.showerror("错误", f"撤回操作失败：{str(e)}")
                # 如果撤回失败，将操作重新加入历史记录
                self.operation_history.append(last_operation)

if __name__ == "__main__":
    root = tk.Tk()
    app = ImageClassifier(root)
    root.mainloop()