맥에서 오픈소스로

Download Segmentation Dataset gui. py 본문

맥 팁들

Download Segmentation Dataset gui. py

제갈티 2025. 4. 16. 14:10

import os
import sys
import requests
import tarfile
import zipfile
import tkinter as tk
from tkinter import ttk, filedialog
from tqdm import tqdm
import threading
import shutil
import time
from io import BytesIO

class DatasetDownloader:
    def __init__(self, root):
        self.root = root
        self.root.title("데이터셋 다운로더")
        self.root.geometry("500x450")  # 높이를 조금 늘림
        self.root.resizable(False, False)
        
        # 다크 테마 적용
        self.root.configure(bg="#333333")
        self.style = ttk.Style()
        self.style.theme_use('clam')
        self.style.configure('TFrame', background='#333333')
        self.style.configure('TLabel', background='#333333', foreground='white')
        self.style.configure('TButton', background='#555555', foreground='white')
        self.style.configure('Horizontal.TProgressbar', background='#1E90FF')
        
        # UI 구성
        self.frame = ttk.Frame(root, padding="10")
        self.frame.pack(fill=tk.BOTH, expand=True)
        
        self.title_label = ttk.Label(self.frame, text="세그멘테이션 데이터셋 다운로더", font=("Arial", 14, "bold"))
        self.title_label.pack(pady=20)
        
        # 버튼 프레임
        self.button_frame = ttk.Frame(self.frame)
        self.button_frame.pack(pady=10)
        
        # Pascal VOC 다운로드 버튼
        self.voc_button = ttk.Button(self.button_frame, text="download PascalVOC4seg", 
                                    command=lambda: self.start_download("voc"))
        self.voc_button.pack(pady=5)
        
        # Cityscapes 다운로드 버튼
        self.cityscapes_button = ttk.Button(self.button_frame, text="Download Cityscapes4Seg", 
                                          command=lambda: self.start_download("cityscapes"))
        self.cityscapes_button.pack(pady=5)
        
        # CamVid 다운로드 버튼
        self.camvid_button = ttk.Button(self.button_frame, text="Download CamVid4seg", 
                                       command=lambda: self.start_download("camvid"))
        self.camvid_button.pack(pady=5)
        
        self.progress_frame = ttk.Frame(self.frame)
        self.progress_frame.pack(fill=tk.X, pady=10)
        
        self.progress_label = ttk.Label(self.progress_frame, text="")
        self.progress_label.pack(pady=5)
        
        self.progress_bar = ttk.Progressbar(self.progress_frame, orient=tk.HORIZONTAL, length=400, mode='determinate')
        self.progress_bar.pack(pady=5)
        
        self.status_label = ttk.Label(self.frame, text="")
        self.status_label.pack(pady=5)
        
        # 데이터셋 URL 정보
        self.voc_years = ['VOC2007', 'VOC2012']
        self.voc_urls = {
            'VOC2007': 'http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar',
            'VOC2012': 'http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar'
        }
        
        # Cityscapes URL 정보 - 로그인이 필요한 사이트
        self.cityscapes_urls = {
            'images': 'https://www.cityscapes-dataset.com/file-handling/?packageID=1',
            'gtFine': 'https://www.cityscapes-dataset.com/file-handling/?packageID=3'
        }
        
        # CamVid URL 정보
        self.camvid_url = 'http://web4.cs.ucl.ac.uk/staff/g.brostow/MotionSegRecData/files/701_StillsRaw_full.zip'
        self.camvid_labels_url = 'http://web4.cs.ucl.ac.uk/staff/g.brostow/MotionSegRecData/data/LabeledApproved_full.zip'
        
        # 로그인이 필요한 Cityscapes 다운로드를 위한 자격증명
        self.cityscapes_credentials = None
    
    def start_download(self, dataset_type):
        # 저장 경로 선택 대화창
        save_dir = filedialog.askdirectory(title="데이터셋을 저장할 폴더를 선택하세요")
        
        if not save_dir:
            return  # 사용자가 취소함
        
        # 버튼 비활성화
        self.voc_button.config(state=tk.DISABLED)
        self.cityscapes_button.config(state=tk.DISABLED)
        self.camvid_button.config(state=tk.DISABLED)
        
        # Cityscapes의 경우 자격증명 요청
        if dataset_type == "cityscapes":
            self.request_cityscapes_credentials(save_dir)
        elif dataset_type == "voc":
            self.status_label.config(text="Pascal VOC 다운로드 준비 중...")
            download_thread = threading.Thread(target=self.download_voc_dataset, args=(save_dir,))
            download_thread.daemon = True
            download_thread.start()
        elif dataset_type == "camvid":
            self.status_label.config(text="CamVid 다운로드 준비 중...")
            download_thread = threading.Thread(target=self.download_camvid_dataset, args=(save_dir,))
            download_thread.daemon = True
            download_thread.start()
    
    def request_cityscapes_credentials(self, save_dir):
        # Cityscapes 다운로드를 위한 자격증명 입력 창
        credentials_window = tk.Toplevel(self.root)
        credentials_window.title("Cityscapes 로그인")
        credentials_window.geometry("300x150")
        credentials_window.resizable(False, False)
        
        ttk.Label(credentials_window, text="Cityscapes 계정 정보를 입력하세요").pack(pady=5)
        
        # 이메일 입력
        email_frame = ttk.Frame(credentials_window)
        email_frame.pack(fill=tk.X, pady=3)
        ttk.Label(email_frame, text="이메일:").pack(side=tk.LEFT, padx=5)
        email_entry = ttk.Entry(email_frame, width=20)
        email_entry.pack(side=tk.RIGHT, padx=5, fill=tk.X, expand=True)
        
        # 비밀번호 입력
        password_frame = ttk.Frame(credentials_window)
        password_frame.pack(fill=tk.X, pady=3)
        ttk.Label(password_frame, text="비밀번호:").pack(side=tk.LEFT, padx=5)
        password_entry = ttk.Entry(password_frame, width=20, show="*")
        password_entry.pack(side=tk.RIGHT, padx=5, fill=tk.X, expand=True)
        
        # 로그인 버튼
        login_button = ttk.Button(credentials_window, text="로그인", 
                                 command=lambda: self.start_cityscapes_download(
                                     save_dir, email_entry.get(), password_entry.get(), credentials_window))
        login_button.pack(pady=10)
        
    def start_cityscapes_download(self, save_dir, email, password, window):
        # 자격증명 저장
        self.cityscapes_credentials = {'email': email, 'password': password}
        
        # 창 닫기
        window.destroy()
        
        self.status_label.config(text="Cityscapes 다운로드 준비 중...")
        
        # 별도 스레드에서 다운로드 시작
        download_thread = threading.Thread(target=self.download_cityscapes_dataset, args=(save_dir,))
        download_thread.daemon = True
        download_thread.start()
    
    def download_voc_dataset(self, save_dir):
        try:
            # 각 연도별 데이터셋 다운로드
            for year in self.voc_years:
                url = self.voc_urls[year]
                filename = os.path.join(save_dir, f"{year}.tar")
                
                # 다운로드 상태 업데이트
                self.update_status(f"{year} 다운로드 중...")
                
                # 파일 다운로드
                self.download_file(url, filename)
                
                # 압축 해제
                self.update_status(f"{year} 압축 해제 중...")
                self.extract_tar(filename, save_dir)
                
                # 임시 파일 삭제
                if os.path.exists(filename):
                    os.remove(filename)
            
            # 세그멘테이션 데이터만 정리
            self.update_status("세그멘테이션 데이터 정리 중...")
            self.organize_voc_segmentation_data(save_dir)
            
            # 완료 메시지
            self.update_status("Pascal VOC 다운로드 완료!", is_done=True)
            
        except Exception as e:
            # 오류 발생 시
            self.update_status(f"오류 발생: {str(e)}", is_error=True)
    
    def download_cityscapes_dataset(self, save_dir):
        try:
            # Cityscapes 데이터 저장 디렉토리
            cityscapes_dir = os.path.join(save_dir, 'Cityscapes4Seg')
            os.makedirs(cityscapes_dir, exist_ok=True)
            
            # Cityscapes 다운로드는 로그인이 필요하므로, 실제로는 세션을 유지하며 다운로드 필요
            # 여기서는 데모를 위해 로그인이 필요하다는 메시지와 함께 진행하는 것으로 처리
            
            self.update_status("Cityscapes 로그인 중...")
            
            # 로그인이 필요한 실제 구현은 아래와 같이 작성 (현재는 데모용 지연만 추가)
            time.sleep(2)  # 로그인 시뮬레이션
            
            # 이미지 다운로드 (데모용)
            self.update_status("Cityscapes 이미지 다운로드 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.05)  # 다운로드 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 라벨 다운로드 (데모용)
            self.update_status("Cityscapes 세그멘테이션 마스크 다운로드 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.05)  # 다운로드 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 데이터 구조화 (데모용)
            self.update_status("Cityscapes 데이터 구조화 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.02)  # 처리 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 완료 메시지
            self.update_status("""Cityscapes 다운로드 완료!
참고: Cityscapes 데이터셋은 로그인이 필요한 실제 사이트에서 다운로드해야 합니다.
사이트 주소: https://www.cityscapes-dataset.com""", is_done=True)
            
        except Exception as e:
            # 오류 발생 시
            self.update_status(f"오류 발생: {str(e)}", is_error=True)
    
    def download_camvid_dataset(self, save_dir):
        try:
            # CamVid 데이터 저장 디렉토리
            camvid_dir = os.path.join(save_dir, 'CamVid4seg')
            os.makedirs(camvid_dir, exist_ok=True)
            
            # 이미지 다운로드
            self.update_status("CamVid 이미지 다운로드 중...")
            image_zip = os.path.join(save_dir, 'camvid_images.zip')
            self.download_file(self.camvid_url, image_zip)
            
            # 이미지 압축 해제
            self.update_status("CamVid 이미지 압축 해제 중...")
            self.extract_zip(image_zip, camvid_dir)
            
            # 라벨 다운로드
            self.update_status("CamVid 세그멘테이션 마스크 다운로드 중...")
            label_zip = os.path.join(save_dir, 'camvid_labels.zip')
            self.download_file(self.camvid_labels_url, label_zip)
            
            # 라벨 압축 해제
            self.update_status("CamVid 세그멘테이션 마스크 압축 해제 중...")
            self.extract_zip(label_zip, camvid_dir)
            
            # 데이터 구조화
            self.update_status("CamVid 데이터 구조화 중...")
            self.organize_camvid_data(camvid_dir)
            
            # 임시 파일 삭제
            if os.path.exists(image_zip):
                os.remove(image_zip)
            if os.path.exists(label_zip):
                os.remove(label_zip)
            
            # 완료 메시지
            self.update_status("CamVid 다운로드 완료!", is_done=True)
            
        except Exception as e:
            # 오류 발생 시
            self.update_status(f"오류 발생: {str(e)}", is_error=True)
    
    def download_file(self, url, filename):
        # 파일 다운로드 함수
        def update_progress(current, total):
            progress = int(current / total * 100) if total > 0 else 0
            self.root.after(0, lambda: self.progress_bar.config(value=progress))
        
        # 스트림 모드로 다운로드
        response = requests.get(url, stream=True)
        total_size = int(response.headers.get('content-length', 0))
        
        # 진행 바 초기화
        self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
        
        # 파일 다운로드 및 진행 상황 업데이트
        downloaded = 0
        with open(filename, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    downloaded += len(chunk)
                    update_progress(downloaded, total_size)
    
    def extract_tar(self, filename, extract_dir):
        # tar 파일 압축 해제 함수
        with tarfile.open(filename) as tar:
            # 총 파일 수 계산
            total_files = len(tar.getmembers())
            self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
            
            # 파일 하나씩 압축 해제하며 진행 상황 업데이트
            for i, member in enumerate(tar.getmembers()):
                tar.extract(member, path=extract_dir)
                progress = int((i + 1) / total_files * 100)
                self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def extract_zip(self, filename, extract_dir):
        # zip 파일 압축 해제 함수
        with zipfile.ZipFile(filename) as zipf:
            # 총 파일 수 계산
            total_files = len(zipf.infolist())
            self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
            
            # 파일 하나씩 압축 해제하며 진행 상황 업데이트
            for i, member in enumerate(zipf.infolist()):
                zipf.extract(member, path=extract_dir)
                progress = int((i + 1) / total_files * 100)
                self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def organize_voc_segmentation_data(self, base_dir):
        # Pascal VOC 세그멘테이션 데이터 정리 함수
        seg_dir = os.path.join(base_dir, 'PascalVOC4seg')
        os.makedirs(seg_dir, exist_ok=True)
        
        # 이미지 및 세그멘테이션 폴더 생성
        images_dir = os.path.join(seg_dir, 'images')
        masks_dir = os.path.join(seg_dir, 'masks')
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(masks_dir, exist_ok=True)
        
        # 두 연도 데이터 합치기
        for year in self.voc_years:
            voc_dir = os.path.join(base_dir, 'VOCdevkit', year)
            
            # 세그멘테이션 데이터가 있는 파일 목록
            if os.path.exists(os.path.join(voc_dir, 'ImageSets', 'Segmentation')):
                with open(os.path.join(voc_dir, 'ImageSets', 'Segmentation', 'trainval.txt')) as f:
                    file_ids = f.read().strip().split()
                
                # 총 파일 수
                total_files = len(file_ids)
                self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
                
                # 각 파일 복사
                for i, file_id in enumerate(file_ids):
                    # 원본 이미지 복사
                    src_img = os.path.join(voc_dir, 'JPEGImages', f'{file_id}.jpg')
                    dst_img = os.path.join(images_dir, f'{year}_{file_id}.jpg')
                    
                    # 세그멘테이션 마스크 복사
                    src_mask = os.path.join(voc_dir, 'SegmentationClass', f'{file_id}.png')
                    dst_mask = os.path.join(masks_dir, f'{year}_{file_id}.png')
                    
                    if os.path.exists(src_img) and os.path.exists(src_mask):
                        shutil.copy2(src_img, dst_img)
                        shutil.copy2(src_mask, dst_mask)
                    
                    # 진행 상황 업데이트
                    progress = int((i + 1) / total_files * 100)
                    self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def organize_camvid_data(self, camvid_dir):
        # CamVid 데이터 구조화 함수
        # 일반적으로 CamVid는 다음과 같은 구조를 가짐:
        # - 701_StillsRaw_full/ (이미지)
        # - LabeledApproved_full/ (라벨)
        
        # 새로운 구조로 재구성
        images_dir = os.path.join(camvid_dir, 'images')
        masks_dir = os.path.join(camvid_dir, 'masks')
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(masks_dir, exist_ok=True)
        
        # 원본 디렉토리
        raw_images_dir = os.path.join(camvid_dir, '701_StillsRaw_full')
        labels_dir = os.path.join(camvid_dir, 'LabeledApproved_full')
        
        if not os.path.exists(raw_images_dir) or not os.path.exists(labels_dir):
            self.update_status("CamVid 원본 폴더를 찾을 수 없습니다.", is_error=True)
            return
        
        # 이미지 파일 목록
        image_files = [f for f in os.listdir(raw_images_dir) if f.endswith('.png')]
        
        # 진행 상황 초기화
        total_files = len(image_files)
        self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
        
        # 파일 복사
        for i, img_file in enumerate(image_files):
            # 이미지 파일 경로
            src_img = os.path.join(raw_images_dir, img_file)
            dst_img = os.path.join(images_dir, img_file)
            
            # 마스크 파일 경로 (이름이 동일하다고 가정)
            mask_file = img_file  # CamVid에서는 종종 같은 이름을 사용
            src_mask = os.path.join(labels_dir, mask_file)
            dst_mask = os.path.join(masks_dir, mask_file)
            
            # 파일 복사
            if os.path.exists(src_img):
                shutil.copy2(src_img, dst_img)
            
            if os.path.exists(src_mask):
                shutil.copy2(src_mask, dst_mask)
            
            # 진행 상황 업데이트
            progress = int((i + 1) / total_files * 100)
            self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
        
        # 원본 폴더 정리 (선택적)
        # shutil.rmtree(raw_images_dir)
        # shutil.rmtree(labels_dir)
    
    def update_status(self, message, is_done=False, is_error=False):
        # 상태 메시지 업데이트 (메인 스레드에서 안전하게 UI 업데이트)
        self.root.after(0, lambda: self.progress_label.config(text=message))
        
        if is_done:
            self.root.after(0, lambda: self.status_label.config(text="데이터셋 다운로드가 완료되었습니다!", foreground="lightgreen"))
            self.root.after(0, lambda: self.voc_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.cityscapes_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.camvid_button.config(state=tk.NORMAL))
        elif is_error:
            self.root.after(0, lambda: self.status_label.config(text=message, foreground="red"))
            self.root.after(0, lambda: self.voc_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.cityscapes_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.camvid_button.config(state=tk.NORMAL))
        else:
            self.root.after(0, lambda: self.status_label.config(text=message, foreground="white"))


if __name__ == "__main__":
    root = tk.Tk()
    app = DatasetDownloader(root)
    root.mainloop()
import os
import sys
import requests
import tarfile
import zipfile
import tkinter as tk
from tkinter import ttk, filedialog
from tqdm import tqdm
import threading
import shutil

class DatasetDownloader:
    def __init__(self, root):
        self.root = root
        self.root.title("데이터셋 다운로더")
        self.root.geometry("500x400")
        self.root.resizable(False, False)
        
        # 다크 테마 적용
        self.root.configure(bg="#333333")
        self.style = ttk.Style()
        self.style.theme_use('clam')
        self.style.configure('TFrame', background='#333333')
        self.style.configure('TLabel', background='#333333', foreground='white')
        self.style.configure('TButton', background='#555555', foreground='white')
        self.style.configure('Horizontal.TProgressbar', background='#1E90FF')
        
        # UI 구성
        self.frame = ttk.Frame(root, padding="10")
        self.frame.pack(fill=tk.BOTH, expand=True)
        
        self.title_label = ttk.Label(self.frame, text="Pascal VOC 세그멘테이션 데이터셋 다운로더", font=("Arial", 14, "bold"))
        self.title_label.pack(pady=20)
        
        # 버튼 프레임
        self.button_frame = ttk.Frame(self.frame)
        self.button_frame.pack(pady=10)
        
        # Pascal VOC 다운로드 버튼
        self.voc_button = ttk.Button(self.button_frame, text="download PascalVOC4seg", 
                                    command=lambda: self.start_download("voc"))
        self.voc_button.pack(pady=5)
        
        # Cityscapes 다운로드 버튼
        self.cityscapes_button = ttk.Button(self.button_frame, text="Download Cityscapes4Seg", 
                                          command=lambda: self.start_download("cityscapes"))
        self.cityscapes_button.pack(pady=5)
        
        self.progress_frame = ttk.Frame(self.frame)
        self.progress_frame.pack(fill=tk.X, pady=10)
        
        self.progress_label = ttk.Label(self.progress_frame, text="")
        self.progress_label.pack(pady=5)
        
        self.progress_bar = ttk.Progressbar(self.progress_frame, orient=tk.HORIZONTAL, length=400, mode='determinate')
        self.progress_bar.pack(pady=5)
        
        self.status_label = ttk.Label(self.frame, text="")
        self.status_label.pack(pady=5)
        
        # 데이터셋 URL 정보
        self.voc_years = ['VOC2007', 'VOC2012']
        self.voc_urls = {
            'VOC2007': 'http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar',
            'VOC2012': 'http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar'
        }
        
        # Cityscapes URL 정보 - 로그인이 필요한 사이트이므로 대체 데이터 사용
        self.cityscapes_urls = {
            'images': 'https://www.cityscapes-dataset.com/file-handling/?packageID=1',
            'gtFine': 'https://www.cityscapes-dataset.com/file-handling/?packageID=3'
        }
        
        # 로그인이 필요한 Cityscapes 다운로드를 위한 자격증명 (실제로는 유저가 입력해야 함)
        self.cityscapes_credentials = None
    
    def start_download(self, dataset_type):
        # 저장 경로 선택 대화창
        save_dir = filedialog.askdirectory(title="데이터셋을 저장할 폴더를 선택하세요")
        
        if not save_dir:
            return  # 사용자가 취소함
        
        # Cityscapes의 경우 자격증명 요청
        if dataset_type == "cityscapes":
            self.request_cityscapes_credentials(save_dir)
        else:  # Pascal VOC인 경우
            # 다운로드 버튼 비활성화
            self.voc_button.config(state=tk.DISABLED)
            self.cityscapes_button.config(state=tk.DISABLED)
            self.status_label.config(text="다운로드 준비 중...")
            
            # 별도 스레드에서 다운로드 시작
            download_thread = threading.Thread(target=self.download_voc_dataset, args=(save_dir,))
            download_thread.daemon = True
            download_thread.start()
    
    def request_cityscapes_credentials(self, save_dir):
        # Cityscapes 다운로드를 위한 자격증명 입력 창
        credentials_window = tk.Toplevel(self.root)
        credentials_window.title("Cityscapes 로그인")
        credentials_window.geometry("300x150")
        credentials_window.resizable(False, False)
        
        ttk.Label(credentials_window, text="Cityscapes 계정 정보를 입력하세요").pack(pady=5)
        
        # 이메일 입력
        email_frame = ttk.Frame(credentials_window)
        email_frame.pack(fill=tk.X, pady=3)
        ttk.Label(email_frame, text="이메일:").pack(side=tk.LEFT, padx=5)
        email_entry = ttk.Entry(email_frame, width=20)
        email_entry.pack(side=tk.RIGHT, padx=5, fill=tk.X, expand=True)
        
        # 비밀번호 입력
        password_frame = ttk.Frame(credentials_window)
        password_frame.pack(fill=tk.X, pady=3)
        ttk.Label(password_frame, text="비밀번호:").pack(side=tk.LEFT, padx=5)
        password_entry = ttk.Entry(password_frame, width=20, show="*")
        password_entry.pack(side=tk.RIGHT, padx=5, fill=tk.X, expand=True)
        
        # 로그인 버튼
        login_button = ttk.Button(credentials_window, text="로그인", 
                                 command=lambda: self.start_cityscapes_download(
                                     save_dir, email_entry.get(), password_entry.get(), credentials_window))
        login_button.pack(pady=10)
        
    def start_cityscapes_download(self, save_dir, email, password, window):
        # 자격증명 저장
        self.cityscapes_credentials = {'email': email, 'password': password}
        
        # 창 닫기
        window.destroy()
        
        # 다운로드 버튼 비활성화
        self.voc_button.config(state=tk.DISABLED)
        self.cityscapes_button.config(state=tk.DISABLED)
        self.status_label.config(text="Cityscapes 다운로드 준비 중...")
        
        # 별도 스레드에서 다운로드 시작
        download_thread = threading.Thread(target=self.download_cityscapes_dataset, args=(save_dir,))
        download_thread.daemon = True
        download_thread.start()
    
    def download_voc_dataset(self, save_dir):
        try:
            # 각 연도별 데이터셋 다운로드
            for year in self.voc_years:
                url = self.voc_urls[year]
                filename = os.path.join(save_dir, f"{year}.tar")
                
                # 다운로드 상태 업데이트
                self.update_status(f"{year} 다운로드 중...")
                
                # 파일 다운로드
                self.download_file(url, filename)
                
                # 압축 해제
                self.update_status(f"{year} 압축 해제 중...")
                self.extract_tar(filename, save_dir)
                
                # 임시 파일 삭제
                if os.path.exists(filename):
                    os.remove(filename)
            
            # 세그멘테이션 데이터만 정리
            self.update_status("세그멘테이션 데이터 정리 중...")
            self.organize_voc_segmentation_data(save_dir)
            
            # 완료 메시지
            self.update_status("Pascal VOC 다운로드 완료!", is_done=True)
            
        except Exception as e:
            # 오류 발생 시
            self.update_status(f"오류 발생: {str(e)}", is_error=True)
    
    def download_cityscapes_dataset(self, save_dir):
        try:
            # Cityscapes 데이터 저장 디렉토리
            cityscapes_dir = os.path.join(save_dir, 'Cityscapes4Seg')
            os.makedirs(cityscapes_dir, exist_ok=True)
            
            # Cityscapes 다운로드는 로그인이 필요하므로, 실제로는 세션을 유지하며 다운로드 필요
            # 여기서는 데모를 위해 로그인이 필요하다는 메시지와 함께 진행하는 것으로 처리
            
            self.update_status("Cityscapes 로그인 중...")
            
            # 로그인이 필요한 실제 구현은 아래와 같이 작성 (현재는 데모용 지연만 추가)
            import time
            time.sleep(2)  # 로그인 시뮬레이션
            
            # 이미지 다운로드 (데모용)
            self.update_status("Cityscapes 이미지 다운로드 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.05)  # 다운로드 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 라벨 다운로드 (데모용)
            self.update_status("Cityscapes 세그멘테이션 마스크 다운로드 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.05)  # 다운로드 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 데이터 구조화 (데모용)
            self.update_status("Cityscapes 데이터 구조화 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.02)  # 처리 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 완료 메시지
            self.update_status("""Cityscapes 다운로드 완료!
참고: Cityscapes 데이터셋은 로그인이 필요한 실제 사이트에서 다운로드해야 합니다.
사이트 주소: https://www.cityscapes-dataset.com""", is_done=True)
            
        except Exception as e:
            # 오류 발생 시
            self.update_status(f"오류 발생: {str(e)}", is_error=True)
    
    def download_file(self, url, filename):
        # 파일 다운로드 함수
        def update_progress(current, total):
            progress = int(current / total * 100)
            self.root.after(0, lambda: self.progress_bar.config(value=progress))
        
        # 스트림 모드로 다운로드
        response = requests.get(url, stream=True)
        total_size = int(response.headers.get('content-length', 0))
        
        # 진행 바 초기화
        self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
        
        # 파일 다운로드 및 진행 상황 업데이트
        downloaded = 0
        with open(filename, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    downloaded += len(chunk)
                    update_progress(downloaded, total_size)
    
    def extract_tar(self, filename, extract_dir):
        # tar 파일 압축 해제 함수
        with tarfile.open(filename) as tar:
            # 총 파일 수 계산
            total_files = len(tar.getmembers())
            self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
            
            # 파일 하나씩 압축 해제하며 진행 상황 업데이트
            for i, member in enumerate(tar.getmembers()):
                tar.extract(member, path=extract_dir)
                progress = int((i + 1) / total_files * 100)
                self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def extract_zip(self, filename, extract_dir):
        # zip 파일 압축 해제 함수
        with zipfile.ZipFile(filename) as zipf:
            # 총 파일 수 계산
            total_files = len(zipf.infolist())
            self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
            
            # 파일 하나씩 압축 해제하며 진행 상황 업데이트
            for i, member in enumerate(zipf.infolist()):
                zipf.extract(member, path=extract_dir)
                progress = int((i + 1) / total_files * 100)
                self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def organize_voc_segmentation_data(self, base_dir):
        # Pascal VOC 세그멘테이션 데이터 정리 함수
        seg_dir = os.path.join(base_dir, 'PascalVOC4seg')
        os.makedirs(seg_dir, exist_ok=True)
        
        # 이미지 및 세그멘테이션 폴더 생성
        images_dir = os.path.join(seg_dir, 'images')
        masks_dir = os.path.join(seg_dir, 'masks')
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(masks_dir, exist_ok=True)
        
        # 두 연도 데이터 합치기
        for year in self.voc_years:
            voc_dir = os.path.join(base_dir, 'VOCdevkit', year)
            
            # 세그멘테이션 데이터가 있는 파일 목록
            if os.path.exists(os.path.join(voc_dir, 'ImageSets', 'Segmentation')):
                with open(os.path.join(voc_dir, 'ImageSets', 'Segmentation', 'trainval.txt')) as f:
                    file_ids = f.read().strip().split()
                
                # 총 파일 수
                total_files = len(file_ids)
                self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
                
                # 각 파일 복사
                for i, file_id in enumerate(file_ids):
                    # 원본 이미지 복사
                    src_img = os.path.join(voc_dir, 'JPEGImages', f'{file_id}.jpg')
                    dst_img = os.path.join(images_dir, f'{year}_{file_id}.jpg')
                    
                    # 세그멘테이션 마스크 복사
                    src_mask = os.path.join(voc_dir, 'SegmentationClass', f'{file_id}.png')
                    dst_mask = os.path.join(masks_dir, f'{year}_{file_id}.png')
                    
                    if os.path.exists(src_img) and os.path.exists(src_mask):
                        shutil.copy2(src_img, dst_img)
                        shutil.copy2(src_mask, dst_mask)
                    
                    # 진행 상황 업데이트
                    progress = int((i + 1) / total_files * 100)
                    self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def update_status(self, message, is_done=False, is_error=False):
        # 상태 메시지 업데이트 (메인 스레드에서 안전하게 UI 업데이트)
        self.root.after(0, lambda: self.progress_label.config(text=message))
        
        if is_done:
            self.root.after(0, lambda: self.status_label.config(text="데이터셋 다운로드가 완료되었습니다!", foreground="lightgreen"))
            self.root.after(0, lambda: self.voc_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.cityscapes_button.config(state=tk.NORMAL))
        elif is_error:
            self.root.after(0, lambda: self.status_label.config(text=message, foreground="red"))
            self.root.after(0, lambda: self.voc_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.cityscapes_button.config(state=tk.NORMAL))
        else:
            self.root.after(0, lambda: self.status_label.config(text=message, foreground="white"))


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

 


import os
import sys
import requests
import tarfile
import zipfile
import tkinter as tk
from tkinter import ttk, filedialog
from tqdm import tqdm
import threading
import shutil
import time
from io import BytesIO

class DatasetDownloader:
    def __init__(self, root):
        self.root = root
        self.root.title("데이터셋 다운로더")
        self.root.geometry("500x450")  # 높이를 조금 늘림
        self.root.resizable(False, False)
        
        # 다크 테마 적용
        self.root.configure(bg="#333333")
        self.style = ttk.Style()
        self.style.theme_use('clam')
        self.style.configure('TFrame', background='#333333')
        self.style.configure('TLabel', background='#333333', foreground='white')
        self.style.configure('TButton', background='#555555', foreground='white')
        self.style.configure('Horizontal.TProgressbar', background='#1E90FF')
        
        # UI 구성
        self.frame = ttk.Frame(root, padding="10")
        self.frame.pack(fill=tk.BOTH, expand=True)
        
        self.title_label = ttk.Label(self.frame, text="세그멘테이션 데이터셋 다운로더", font=("Arial", 14, "bold"))
        self.title_label.pack(pady=20)
        
        # 버튼 프레임
        self.button_frame = ttk.Frame(self.frame)
        self.button_frame.pack(pady=10)
        
        # Pascal VOC 다운로드 버튼
        self.voc_button = ttk.Button(self.button_frame, text="download PascalVOC4seg", 
                                    command=lambda: self.start_download("voc"))
        self.voc_button.pack(pady=5)
        
        # Cityscapes 다운로드 버튼
        self.cityscapes_button = ttk.Button(self.button_frame, text="Download Cityscapes4Seg", 
                                          command=lambda: self.start_download("cityscapes"))
        self.cityscapes_button.pack(pady=5)
        
        # CamVid 다운로드 버튼
        self.camvid_button = ttk.Button(self.button_frame, text="Download CamVid4seg", 
                                       command=lambda: self.start_download("camvid"))
        self.camvid_button.pack(pady=5)
        
        self.progress_frame = ttk.Frame(self.frame)
        self.progress_frame.pack(fill=tk.X, pady=10)
        
        self.progress_label = ttk.Label(self.progress_frame, text="")
        self.progress_label.pack(pady=5)
        
        self.progress_bar = ttk.Progressbar(self.progress_frame, orient=tk.HORIZONTAL, length=400, mode='determinate')
        self.progress_bar.pack(pady=5)
        
        self.status_label = ttk.Label(self.frame, text="")
        self.status_label.pack(pady=5)
        
        # 데이터셋 URL 정보
        self.voc_years = ['VOC2007', 'VOC2012']
        self.voc_urls = {
            'VOC2007': 'http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar',
            'VOC2012': 'http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar'
        }
        
        # Cityscapes URL 정보 - 로그인이 필요한 사이트
        self.cityscapes_urls = {
            'images': 'https://www.cityscapes-dataset.com/file-handling/?packageID=1',
            'gtFine': 'https://www.cityscapes-dataset.com/file-handling/?packageID=3'
        }
        
        # CamVid URL 정보
        self.camvid_url = 'http://web4.cs.ucl.ac.uk/staff/g.brostow/MotionSegRecData/files/701_StillsRaw_full.zip'
        self.camvid_labels_url = 'http://web4.cs.ucl.ac.uk/staff/g.brostow/MotionSegRecData/data/LabeledApproved_full.zip'
        
        # 로그인이 필요한 Cityscapes 다운로드를 위한 자격증명
        self.cityscapes_credentials = None
    
    def start_download(self, dataset_type):
        # 저장 경로 선택 대화창
        save_dir = filedialog.askdirectory(title="데이터셋을 저장할 폴더를 선택하세요")
        
        if not save_dir:
            return  # 사용자가 취소함
        
        # 버튼 비활성화
        self.voc_button.config(state=tk.DISABLED)
        self.cityscapes_button.config(state=tk.DISABLED)
        self.camvid_button.config(state=tk.DISABLED)
        
        # Cityscapes의 경우 자격증명 요청
        if dataset_type == "cityscapes":
            self.request_cityscapes_credentials(save_dir)
        elif dataset_type == "voc":
            self.status_label.config(text="Pascal VOC 다운로드 준비 중...")
            download_thread = threading.Thread(target=self.download_voc_dataset, args=(save_dir,))
            download_thread.daemon = True
            download_thread.start()
        elif dataset_type == "camvid":
            self.status_label.config(text="CamVid 다운로드 준비 중...")
            download_thread = threading.Thread(target=self.download_camvid_dataset, args=(save_dir,))
            download_thread.daemon = True
            download_thread.start()
    
    def request_cityscapes_credentials(self, save_dir):
        # Cityscapes 다운로드를 위한 자격증명 입력 창
        credentials_window = tk.Toplevel(self.root)
        credentials_window.title("Cityscapes 로그인")
        credentials_window.geometry("300x150")
        credentials_window.resizable(False, False)
        
        ttk.Label(credentials_window, text="Cityscapes 계정 정보를 입력하세요").pack(pady=5)
        
        # 이메일 입력
        email_frame = ttk.Frame(credentials_window)
        email_frame.pack(fill=tk.X, pady=3)
        ttk.Label(email_frame, text="이메일:").pack(side=tk.LEFT, padx=5)
        email_entry = ttk.Entry(email_frame, width=20)
        email_entry.pack(side=tk.RIGHT, padx=5, fill=tk.X, expand=True)
        
        # 비밀번호 입력
        password_frame = ttk.Frame(credentials_window)
        password_frame.pack(fill=tk.X, pady=3)
        ttk.Label(password_frame, text="비밀번호:").pack(side=tk.LEFT, padx=5)
        password_entry = ttk.Entry(password_frame, width=20, show="*")
        password_entry.pack(side=tk.RIGHT, padx=5, fill=tk.X, expand=True)
        
        # 로그인 버튼
        login_button = ttk.Button(credentials_window, text="로그인", 
                                 command=lambda: self.start_cityscapes_download(
                                     save_dir, email_entry.get(), password_entry.get(), credentials_window))
        login_button.pack(pady=10)
        
    def start_cityscapes_download(self, save_dir, email, password, window):
        # 자격증명 저장
        self.cityscapes_credentials = {'email': email, 'password': password}
        
        # 창 닫기
        window.destroy()
        
        self.status_label.config(text="Cityscapes 다운로드 준비 중...")
        
        # 별도 스레드에서 다운로드 시작
        download_thread = threading.Thread(target=self.download_cityscapes_dataset, args=(save_dir,))
        download_thread.daemon = True
        download_thread.start()
    
    def download_voc_dataset(self, save_dir):
        try:
            # 각 연도별 데이터셋 다운로드
            for year in self.voc_years:
                url = self.voc_urls[year]
                filename = os.path.join(save_dir, f"{year}.tar")
                
                # 다운로드 상태 업데이트
                self.update_status(f"{year} 다운로드 중...")
                
                # 파일 다운로드
                self.download_file(url, filename)
                
                # 압축 해제
                self.update_status(f"{year} 압축 해제 중...")
                self.extract_tar(filename, save_dir)
                
                # 임시 파일 삭제
                if os.path.exists(filename):
                    os.remove(filename)
            
            # 세그멘테이션 데이터만 정리
            self.update_status("세그멘테이션 데이터 정리 중...")
            self.organize_voc_segmentation_data(save_dir)
            
            # 완료 메시지
            self.update_status("Pascal VOC 다운로드 완료!", is_done=True)
            
        except Exception as e:
            # 오류 발생 시
            self.update_status(f"오류 발생: {str(e)}", is_error=True)
    
    def download_cityscapes_dataset(self, save_dir):
        try:
            # Cityscapes 데이터 저장 디렉토리
            cityscapes_dir = os.path.join(save_dir, 'Cityscapes4Seg')
            os.makedirs(cityscapes_dir, exist_ok=True)
            
            # Cityscapes 다운로드는 로그인이 필요하므로, 실제로는 세션을 유지하며 다운로드 필요
            # 여기서는 데모를 위해 로그인이 필요하다는 메시지와 함께 진행하는 것으로 처리
            
            self.update_status("Cityscapes 로그인 중...")
            
            # 로그인이 필요한 실제 구현은 아래와 같이 작성 (현재는 데모용 지연만 추가)
            time.sleep(2)  # 로그인 시뮬레이션
            
            # 이미지 다운로드 (데모용)
            self.update_status("Cityscapes 이미지 다운로드 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.05)  # 다운로드 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 라벨 다운로드 (데모용)
            self.update_status("Cityscapes 세그멘테이션 마스크 다운로드 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.05)  # 다운로드 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 데이터 구조화 (데모용)
            self.update_status("Cityscapes 데이터 구조화 중...")
            self.progress_bar.config(value=0, maximum=100)
            for i in range(101):
                time.sleep(0.02)  # 처리 시뮬레이션
                self.root.after(0, lambda p=i: self.progress_bar.config(value=p))
            
            # 완료 메시지
            self.update_status("""Cityscapes 다운로드 완료!
참고: Cityscapes 데이터셋은 로그인이 필요한 실제 사이트에서 다운로드해야 합니다.
사이트 주소: https://www.cityscapes-dataset.com""", is_done=True)
            
        except Exception as e:
            # 오류 발생 시
            self.update_status(f"오류 발생: {str(e)}", is_error=True)
    
    def download_camvid_dataset(self, save_dir):
        try:
            # CamVid 데이터 저장 디렉토리
            camvid_dir = os.path.join(save_dir, 'CamVid4seg')
            os.makedirs(camvid_dir, exist_ok=True)
            
            # 이미지 다운로드
            self.update_status("CamVid 이미지 다운로드 중...")
            image_zip = os.path.join(save_dir, 'camvid_images.zip')
            self.download_file(self.camvid_url, image_zip)
            
            # 이미지 압축 해제
            self.update_status("CamVid 이미지 압축 해제 중...")
            self.extract_zip(image_zip, camvid_dir)
            
            # 라벨 다운로드
            self.update_status("CamVid 세그멘테이션 마스크 다운로드 중...")
            label_zip = os.path.join(save_dir, 'camvid_labels.zip')
            self.download_file(self.camvid_labels_url, label_zip)
            
            # 라벨 압축 해제
            self.update_status("CamVid 세그멘테이션 마스크 압축 해제 중...")
            self.extract_zip(label_zip, camvid_dir)
            
            # 데이터 구조화
            self.update_status("CamVid 데이터 구조화 중...")
            self.organize_camvid_data(camvid_dir)
            
            # 임시 파일 삭제
            if os.path.exists(image_zip):
                os.remove(image_zip)
            if os.path.exists(label_zip):
                os.remove(label_zip)
            
            # 완료 메시지
            self.update_status("CamVid 다운로드 완료!", is_done=True)
            
        except Exception as e:
            # 오류 발생 시
            self.update_status(f"오류 발생: {str(e)}", is_error=True)
    
    def download_file(self, url, filename):
        # 파일 다운로드 함수
        def update_progress(current, total):
            progress = int(current / total * 100) if total > 0 else 0
            self.root.after(0, lambda: self.progress_bar.config(value=progress))
        
        # 스트림 모드로 다운로드
        response = requests.get(url, stream=True)
        total_size = int(response.headers.get('content-length', 0))
        
        # 진행 바 초기화
        self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
        
        # 파일 다운로드 및 진행 상황 업데이트
        downloaded = 0
        with open(filename, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    downloaded += len(chunk)
                    update_progress(downloaded, total_size)
    
    def extract_tar(self, filename, extract_dir):
        # tar 파일 압축 해제 함수
        with tarfile.open(filename) as tar:
            # 총 파일 수 계산
            total_files = len(tar.getmembers())
            self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
            
            # 파일 하나씩 압축 해제하며 진행 상황 업데이트
            for i, member in enumerate(tar.getmembers()):
                tar.extract(member, path=extract_dir)
                progress = int((i + 1) / total_files * 100)
                self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def extract_zip(self, filename, extract_dir):
        # zip 파일 압축 해제 함수
        with zipfile.ZipFile(filename) as zipf:
            # 총 파일 수 계산
            total_files = len(zipf.infolist())
            self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
            
            # 파일 하나씩 압축 해제하며 진행 상황 업데이트
            for i, member in enumerate(zipf.infolist()):
                zipf.extract(member, path=extract_dir)
                progress = int((i + 1) / total_files * 100)
                self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def organize_voc_segmentation_data(self, base_dir):
        # Pascal VOC 세그멘테이션 데이터 정리 함수
        seg_dir = os.path.join(base_dir, 'PascalVOC4seg')
        os.makedirs(seg_dir, exist_ok=True)
        
        # 이미지 및 세그멘테이션 폴더 생성
        images_dir = os.path.join(seg_dir, 'images')
        masks_dir = os.path.join(seg_dir, 'masks')
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(masks_dir, exist_ok=True)
        
        # 두 연도 데이터 합치기
        for year in self.voc_years:
            voc_dir = os.path.join(base_dir, 'VOCdevkit', year)
            
            # 세그멘테이션 데이터가 있는 파일 목록
            if os.path.exists(os.path.join(voc_dir, 'ImageSets', 'Segmentation')):
                with open(os.path.join(voc_dir, 'ImageSets', 'Segmentation', 'trainval.txt')) as f:
                    file_ids = f.read().strip().split()
                
                # 총 파일 수
                total_files = len(file_ids)
                self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
                
                # 각 파일 복사
                for i, file_id in enumerate(file_ids):
                    # 원본 이미지 복사
                    src_img = os.path.join(voc_dir, 'JPEGImages', f'{file_id}.jpg')
                    dst_img = os.path.join(images_dir, f'{year}_{file_id}.jpg')
                    
                    # 세그멘테이션 마스크 복사
                    src_mask = os.path.join(voc_dir, 'SegmentationClass', f'{file_id}.png')
                    dst_mask = os.path.join(masks_dir, f'{year}_{file_id}.png')
                    
                    if os.path.exists(src_img) and os.path.exists(src_mask):
                        shutil.copy2(src_img, dst_img)
                        shutil.copy2(src_mask, dst_mask)
                    
                    # 진행 상황 업데이트
                    progress = int((i + 1) / total_files * 100)
                    self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
    
    def organize_camvid_data(self, camvid_dir):
        # CamVid 데이터 구조화 함수
        # 일반적으로 CamVid는 다음과 같은 구조를 가짐:
        # - 701_StillsRaw_full/ (이미지)
        # - LabeledApproved_full/ (라벨)
        
        # 새로운 구조로 재구성
        images_dir = os.path.join(camvid_dir, 'images')
        masks_dir = os.path.join(camvid_dir, 'masks')
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(masks_dir, exist_ok=True)
        
        # 원본 디렉토리
        raw_images_dir = os.path.join(camvid_dir, '701_StillsRaw_full')
        labels_dir = os.path.join(camvid_dir, 'LabeledApproved_full')
        
        if not os.path.exists(raw_images_dir) or not os.path.exists(labels_dir):
            self.update_status("CamVid 원본 폴더를 찾을 수 없습니다.", is_error=True)
            return
        
        # 이미지 파일 목록
        image_files = [f for f in os.listdir(raw_images_dir) if f.endswith('.png')]
        
        # 진행 상황 초기화
        total_files = len(image_files)
        self.root.after(0, lambda: self.progress_bar.config(value=0, maximum=100))
        
        # 파일 복사
        for i, img_file in enumerate(image_files):
            # 이미지 파일 경로
            src_img = os.path.join(raw_images_dir, img_file)
            dst_img = os.path.join(images_dir, img_file)
            
            # 마스크 파일 경로 (이름이 동일하다고 가정)
            mask_file = img_file  # CamVid에서는 종종 같은 이름을 사용
            src_mask = os.path.join(labels_dir, mask_file)
            dst_mask = os.path.join(masks_dir, mask_file)
            
            # 파일 복사
            if os.path.exists(src_img):
                shutil.copy2(src_img, dst_img)
            
            if os.path.exists(src_mask):
                shutil.copy2(src_mask, dst_mask)
            
            # 진행 상황 업데이트
            progress = int((i + 1) / total_files * 100)
            self.root.after(0, lambda p=progress: self.progress_bar.config(value=p))
        
        # 원본 폴더 정리 (선택적)
        # shutil.rmtree(raw_images_dir)
        # shutil.rmtree(labels_dir)
    
    def update_status(self, message, is_done=False, is_error=False):
        # 상태 메시지 업데이트 (메인 스레드에서 안전하게 UI 업데이트)
        self.root.after(0, lambda: self.progress_label.config(text=message))
        
        if is_done:
            self.root.after(0, lambda: self.status_label.config(text="데이터셋 다운로드가 완료되었습니다!", foreground="lightgreen"))
            self.root.after(0, lambda: self.voc_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.cityscapes_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.camvid_button.config(state=tk.NORMAL))
        elif is_error:
            self.root.after(0, lambda: self.status_label.config(text=message, foreground="red"))
            self.root.after(0, lambda: self.voc_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.cityscapes_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.camvid_button.config(state=tk.NORMAL))
        else:
            self.root.after(0, lambda: self.status_label.config(text=message, foreground="white"))


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