R360-UAVmodelTool/target_based_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 TargetBasedClassifier:
    def __init__(self, root):
        self.root = root
        self.root.title("基于目标的图片分类器")
        self.root.geometry("1400x900")
        
        # 当前工作目录（包含original、targets、annotated三个文件夹）
        self.work_directory = None
        # 三个子目录路径
        self.original_dir = None
        self.targets_dir = None
        self.annotated_dir = None
        
        # 图片列表和索引
        self.image_list = []  # 存储original文件夹中的图片路径
        self.current_index = 0
        
        # 当前显示的图片信息
        self.current_original_path = None
        self.current_targets = []  # 当前图片对应的目标切片列表
        
        # 图片处理参数
        self.zoom_factor = 1.0
        self.rotation_angle = 0
        self.brightness_factor = 1.0
        
        # 操作历史记录
        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_work_directory).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("manned")).pack(pady=5)
        ttk.Button(control_frame, text="其他 (4)", command=lambda: self.classify_image("other")).pack(pady=5)
        ttk.Button(control_frame, text="无目标 (5)", 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="撤回操作 (6)", 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.Button(process_frame, text="重置图片 (R)", command=self.reset_image).pack(pady=5)
        
        # 右侧显示区域
        display_frame = ttk.Frame(main_frame)
        display_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        
        # 上方：标注图片显示
        original_frame = ttk.LabelFrame(display_frame, text="标注图片（带检测框）")
        original_frame.pack(fill=tk.BOTH, expand=True, pady=(0, 5))
        
        self.original_label = ttk.Label(original_frame)
        self.original_label.pack(fill=tk.BOTH, expand=True)
        
        # 下方：目标切片显示（仅作参考）
        targets_frame = ttk.LabelFrame(display_frame, text="检测到的目标（仅供参考）")
        targets_frame.pack(fill=tk.X, pady=(5, 0))
        
        # 创建可滚动的目标显示区域
        targets_canvas = tk.Canvas(targets_frame, height=200)
        targets_scrollbar = ttk.Scrollbar(targets_frame, orient="horizontal", command=targets_canvas.xview)
        self.targets_scroll_frame = ttk.Frame(targets_canvas)
        
        self.targets_scroll_frame.bind(
            "<Configure>",
            lambda e: targets_canvas.configure(scrollregion=targets_canvas.bbox("all"))
        )
        
        targets_canvas.create_window((0, 0), window=self.targets_scroll_frame, anchor="nw")
        targets_canvas.configure(xscrollcommand=targets_scrollbar.set)
        
        targets_canvas.pack(side="top", fill="both", expand=True)
        targets_scrollbar.pack(side="bottom", fill="x")
        
        # 状态标签
        self.status_label = ttk.Label(self.root, text="请选择包含original、targets、annotated文件夹的工作目录（显示标注图片，分类原图）")
        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_work_directory())
        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("manned"))
        self.root.bind('4', lambda e: self.classify_image("other"))
        self.root.bind('5', lambda e: self.classify_image("none"))
        self.root.bind('6', lambda e: self.undo_operation())
        self.root.bind('r', lambda e: self.reset_image())
        
    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_work_directory(self):
        """选择包含original、targets、annotated三个文件夹的工作目录"""
        folder_path = filedialog.askdirectory(title="选择包含original、targets、annotated文件夹的工作目录")
        if folder_path:
            # 检查是否包含必要的子文件夹
            original_path = os.path.join(folder_path, 'original')
            targets_path = os.path.join(folder_path, 'targets')
            annotated_path = os.path.join(folder_path, 'annotated')
            
            if not all(os.path.exists(path) for path in [original_path, targets_path, annotated_path]):
                messagebox.showerror("错误", "所选目录必须包含original、targets、annotated三个文件夹")
                return
                
            self.work_directory = folder_path
            self.original_dir = original_path
            self.targets_dir = targets_path
            self.annotated_dir = annotated_path
            
            self.load_images()
            
    def load_images(self):
        """加载original文件夹中的图片"""
        self.image_list = []
        
        if not self.original_dir:
            return
            
        for file in os.listdir(self.original_dir):
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                self.image_list.append(os.path.join(self.original_dir, 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("提示", "original文件夹中没有图片")
            
    def find_target_images(self, original_filename):
        """查找与原图对应的目标切片"""
        if not self.targets_dir:
            return []
            
        base_name = os.path.splitext(original_filename)[0]
        target_files = []
        
        for file in os.listdir(self.targets_dir):
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                if file.startswith(base_name + '_'):
                    target_files.append(os.path.join(self.targets_dir, file))
                    
        return sorted(target_files)
        
    def find_annotated_image(self, original_filename):
        """查找与原图对应的标注图片（带检测框的图片）"""
        if not self.annotated_dir:
            return None
            
        base_name = os.path.splitext(original_filename)[0]
        
        # 尝试多种可能的文件扩展名
        for ext in ['.png', '.jpg', '.jpeg', '.bmp', '.gif']:
            annotated_path = os.path.join(self.annotated_dir, base_name + ext)
            if os.path.exists(annotated_path):
                return annotated_path
                
        return None
        
    def show_current_image(self):
        """显示当前图片和对应的目标切片"""
        if 0 <= self.current_index < len(self.image_list):
            self.current_original_path = self.image_list[self.current_index]
            
            # 显示标注图片（带检测框）
            self.display_original_image()
            
            # 显示目标切片
            original_filename = os.path.basename(self.current_original_path)
            self.current_targets = self.find_target_images(original_filename)
            self.display_target_images()
            
            # 更新状态信息
            self.status_label.config(text=f"图片 {self.current_index + 1}/{len(self.image_list)} | 检测到 {len(self.current_targets)} 个目标")
            
            # 显示图片信息
            if os.path.exists(self.current_original_path):
                file_size = os.path.getsize(self.current_original_path) / 1024
                image = Image.open(self.current_original_path)
                width, height = image.size
                self.info_label.config(text=f"文件名: {original_filename} | 尺寸: {width}x{height} | 大小: {file_size:.1f}KB")
            
    def display_original_image(self):
        """显示标注图片（带检测框），如果没有标注图片则显示原图"""
        if not self.current_original_path or not os.path.exists(self.current_original_path):
            self.original_label.config(image='', text="图片不存在")
            return
            
        try:
            # 首先尝试查找标注图片
            original_filename = os.path.basename(self.current_original_path)
            annotated_path = self.find_annotated_image(original_filename)
            
            # 如果有标注图片就显示标注图片，否则显示原图
            if annotated_path and os.path.exists(annotated_path):
                image = Image.open(annotated_path)
                display_text = "（显示标注图片）"
            else:
                image = Image.open(self.current_original_path)
                display_text = "（无标注图片，显示原图）"
            
            # 应用亮度调整
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(self.brightness_factor)
            
            # 应用旋转
            image = image.rotate(self.rotation_angle, expand=True)
            
            # 计算调整后的尺寸，保持宽高比
            width, height = image.size
            max_size = (600, 400)
            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.original_label.config(image=photo, text="")
            self.original_label.image = photo  # 保持引用
            
            # 更新框架标题以显示当前显示的图片类型
            parent_frame = self.original_label.master
            if hasattr(parent_frame, 'config'):
                parent_frame.config(text=f"标注图片（带检测框）{display_text}")
            
        except Exception as e:
            self.original_label.config(image='', text=f"图片加载失败: {str(e)}")
            
    def display_target_images(self):
        """显示目标切片"""
        # 清空之前的目标显示
        for widget in self.targets_scroll_frame.winfo_children():
            widget.destroy()
            
        if not self.current_targets:
            no_target_label = ttk.Label(self.targets_scroll_frame, text="未检测到目标")
            no_target_label.pack(side=tk.LEFT, padx=10)
            return
            
        # 显示每个目标切片
        for i, target_path in enumerate(self.current_targets):
            try:
                # 加载目标图片
                target_image = Image.open(target_path)
                
                # 调整大小以适应显示
                target_image.thumbnail((150, 150), Image.Resampling.LANCZOS)
                target_photo = ImageTk.PhotoImage(target_image)
                
                # 创建目标显示框架
                target_frame = ttk.Frame(self.targets_scroll_frame)
                target_frame.pack(side=tk.LEFT, padx=5, pady=5)
                
                # 显示目标图片
                target_label = ttk.Label(target_frame, image=target_photo)
                target_label.image = target_photo  # 保持引用
                target_label.pack()
                
                # 显示目标文件名
                filename_label = ttk.Label(target_frame, text=os.path.basename(target_path), 
                                         font=('Arial', 8))
                filename_label.pack()
                
            except Exception as e:
                error_label = ttk.Label(self.targets_scroll_frame, text=f"目标{i+1}加载失败")
                error_label.pack(side=tk.LEFT, padx=5)
                
    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_original_path:
            return
            
        # 创建分类文件夹
        category_folders = {
            "drone": "drone",
            "bird": "bird",
            "manned": "manned",
            "other": "others",
            "none": "none"
        }
        
        target_folder = os.path.join(self.work_directory, category_folders[category])
        os.makedirs(target_folder, exist_ok=True)
        
        # 只移动原图
        original_filename = os.path.basename(self.current_original_path)
        target_path = os.path.join(target_folder, original_filename)
        
        try:
            # 记录操作历史（在移动之前）
            operation_record = {
                'action': 'classify',
                'category': category,
                'original_source': self.current_original_path,
                'original_target': target_path,
                'index': self.current_index,
                'original_filename': original_filename
            }
            
            # 执行文件移动（仅移动原图）
            if os.path.exists(self.current_original_path):
                shutil.move(self.current_original_path, target_path)
                operation_record['moved'] = True
            else:
                operation_record['moved'] = False
                
            # 添加到操作历史
            self.operation_history.append(operation_record)
            
            # 限制历史记录数量
            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.original_label.config(image='', text="所有图片已分类完成")
                # 清空目标显示
                for widget in self.targets_scroll_frame.winfo_children():
                    widget.destroy()
                self.status_label.config(text="所有图片已分类完成")
                self.info_label.config(text="")
                
        except Exception as e:
            messagebox.showerror("错误", f"分类操作失败：{str(e)}")
            
    def undo_operation(self):
        """撤回最后一次分类操作"""
        if not self.operation_history:
            messagebox.showinfo("提示", "没有可撤回的操作")
            return
            
        # 获取最后一次操作
        last_operation = self.operation_history.pop()
        
        if last_operation['action'] == 'classify' and last_operation.get('moved', False):
            try:
                # 将原图移回原位置
                if os.path.exists(last_operation['original_target']):
                    shutil.move(last_operation['original_target'], last_operation['original_source'])
                        
                # 重新加载图片列表
                self.load_images()
                
                # 尝试定位到撤回的图片
                original_path = last_operation['original_source']
                        
                if original_path and original_path in self.image_list:
                    self.current_index = self.image_list.index(original_path)
                else:
                    self.current_index = 0
                    
                if self.image_list:
                    self.show_current_image()
                    messagebox.showinfo("成功", f"已撤回对 '{last_operation['original_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 = TargetBasedClassifier(root)
    root.mainloop()