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()