Colab带进度条的拷贝和解压

5 个月前

Untitled

愿你的生活常温暖日子总是温柔又闪光

背景

在Colab中常常要进行大量数据的拷贝工作。巨量小文件对于GoogleDrive而言是会要命的，所以一般的数据集都要压缩成单一压缩包在GoogleDrive和Colab之间传输。而linux系统自带的cp和unzip命令都是不带进度条的。cp命令一执行，什么东西也没有显示，不知道文件多大，不知道拷贝速度如何，不知道拷贝了多少，不知道拷贝还要多久，直接就是一个抓瞎的状态。偶尔大文件拷贝还会失败，基本就是听天由命。unzip命令则是每解压缩一个文件就会打印一条输出，在解压几十万个文件的时候画面太美。为解决上述问题，我被迫自己写了两个可以显示进度的方法 copy_file 和 unzip_file 。目前只针对单文件，暂时没有递归处理文件夹的打算。

效果预览

copy_file.py

import shutil
import os
from tqdm import tqdm
def get_formal_size(byte_size):
    if byte_size < 1024:
        return f"{byte_size} bytes"
    elif byte_size < pow(1024, 2):
        return f"{round(byte_size / 1024, 2)} KB"
    elif byte_size < pow(1024, 3):
        return f"{round(byte_size / (pow(1024, 2)), 2)} MB"
    elif byte_size < pow(1024, 4):
        return f"{round(byte_size / (pow(1024, 3)), 2)} GB"
def copy_with_progress(src_file, dst_file, buffer_size=None):
    # Get the size of the source file (bytes)
    src_total_size = os.path.getsize(src_file)
    # Create a custom file object that updates a progress bar
    class ProgressFileObject:
        def __init__(self, file_object, total_size):
            self.file_object = file_object
            self.total_size = total_size
            self.bytes_transferred = 0
            self.update_counter = 0
            self.cumulated = 0
        def write(self, data):
            self.file_object.write(data)
            self.bytes_transferred += len(data)
            pbar.update(len(data))
    # Copy the file using the custom file object
    with open(src_file, "rb") as f_src:
        with open(dst_file, "wb") as f_dst:
            print('Copy from ' + src_file + ' to ' + dst_file)
            print('File size: ' + get_formal_size(src_total_size))
            with tqdm(total=src_total_size) as pbar:
                pbar.set_description('Progress (bytes)')
                progress_file_object = ProgressFileObject(f_dst, src_total_size)
                if buffer_size is not None:
                    shutil.copyfileobj(f_src, progress_file_object, buffer_size)
                else:
                    shutil.copyfileobj(f_src, progress_file_object)

unzip_file.py

from tqdm import tqdm
import zipfile
def unzip_with_progress(zip_path, extract_path):
    # 打开zip文件并获取其中的文件列表
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        file_list = zip_ref.namelist()