Move files to Clothes-Category-Classifier

.gitignore

@@ -0,0 +1,31 @@
+# Ignore Python bytecode files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Ignore virtual environments
+env/
+venv/
+myenv/
+
+# Ignore Jupyter notebook checkpoints
+.ipynb_checkpoints
+
+# Ignore model files that are tracked by Git LFS
+netBest.pth
+
+# Ignore test images (if you don't want to include them in your repo)
+test_images/
+
+# Ignore log files and temporary files
+*.log
+*.tmp
+*.bak
+
+# Ignore system files
+.DS_Store
+Thumbs.db
+
+# Ignore any other temporary or backup files
+*.swp
+*.swo

main.py

@@ -0,0 +1,100 @@
+import gradio as gr
+import os
+from PIL import Image
+import torch
+import torchvision
+from torchvision import transforms, datasets  # Correct import for torchvision
+import sys
+
+
+# Add the 'src' folder to the system path for module imports
+sys.path.append('./src')
+
+# Import from the 'src' folder
+from model_jigsaw import mae_vit_small_patch16  # Import directly from src folder
+
+
+# Static Variables
+MODEL_PATH = "model/netBest.pth"
+TEST_IMAGE_FOLDER = "test_images/"
+
+# Class Mapping from Clothing1M dataset
+class_names = {
+    0: "T-shirt", 
+    1: "Shirt", 
+    2: "Knitwear", 
+    3: "Chiffon", 
+    4: "Sweater", 
+    5: "Hoodie", 
+    6: "Windbreaker", 
+    7: "Jacket", 
+    8: "Down Coat", 
+    9: "Suit", 
+    10: "Shawl", 
+    11: "Dress", 
+    12: "Vest", 
+    13: "Nightwear"
+}
+
+# Image Preprocessing
+def preprocess_image(image):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    
+    image = image.convert('RGB')
+    image = transform(image).unsqueeze(0)
+    return image
+
+# Load Model
+model = mae_vit_small_patch16(nb_cls=14)
+model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu"))['net'])
+model.eval()
+
+# Load dataset mapping
+val_dataset = torchvision.datasets.ImageFolder(root=TEST_IMAGE_FOLDER)
+idx_to_class = {v: k for k, v in val_dataset.class_to_idx.items()}
+
+def predict_single_image(image):
+    """Predicts the class of an image using the exact logic from eval.py"""
+    image = preprocess_image(image)
+
+    with torch.no_grad():
+        outputs = model.forward_cls(image)
+
+    _, predicted_class = torch.max(outputs, 1)
+
+    # Get final class name
+    mapped_class_index = idx_to_class[predicted_class.item()]
+    final_class_name = class_names[int(mapped_class_index)]
+
+    return final_class_name
+
+# Get all images from subfolders (recursively)
+def get_image_paths(test_image_folder):
+    image_paths = []
+    for root, dirs, files in os.walk(test_image_folder):
+        for file in files:
+            if file.endswith((".jpg", ".png")):
+                image_paths.append(os.path.join(root, file))
+    return image_paths
+
+test_image_files = get_image_paths(TEST_IMAGE_FOLDER)
+
+def load_test_image(selected_image):
+    """Loads the selected test image from test_images folder"""
+    return Image.open(selected_image)
+
+# Create Gradio Interface with dynamic examples from test_images folder
+demo = gr.Interface(
+    fn=predict_single_image,
+    inputs=gr.Image(type="pil", label="Upload an image"),
+    outputs=gr.Textbox(label="Predicted Category"),
+    title="Clothes Category Classifier",
+    description="Upload an image to classify its clothing category.",
+    examples=test_image_files  # Use the correct paths for example images
+)
+
+demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+gradio
+Pillow
+einops
+timm

src/datasets.py

@@ -0,0 +1,91 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+
+import os
+import PIL
+
+from torchvision import datasets, transforms
+
+from timm.data import create_transform
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+
+def build_dataset(is_train, args) :
+    transform = build_transform(is_train, args)
+
+    if args.data_set == 'CIFAR10' :
+        root = os.path.join(args.data_path, 'train' if is_train else 'val')
+        nb_cls = 10
+    elif args.data_set == 'CIFAR100' :
+        root = os.path.join(args.data_path, 'train' if is_train else 'val')
+        nb_cls = 100
+    elif args.data_set == 'Animal10N' :
+        root = os.path.join(args.data_path, 'train' if is_train else 'test')
+        nb_cls = 10
+    elif args.data_set == 'Clothing1M' :
+        # we use a randomly selected balanced training subset
+        root = os.path.join(args.data_path, 'noisy_rand_subtrain' if is_train else 'clean_val')
+        nb_cls = 14
+    elif args.data_set == 'Food101N' :
+        root = os.path.join(args.data_path, 'train' if is_train else 'test')
+        nb_cls = 101
+
+    dataset = datasets.ImageFolder(root, transform=transform)
+
+    print(dataset)
+
+    return dataset, nb_cls
+
+
+def build_transform(is_train, args) :
+    if args.data_set == 'CIFAR10' or args.data_set == 'CIFAR100' :
+        mean = (0.4914, 0.4822, 0.4465)
+        std = (0.2023, 0.1994, 0.2010)
+    else :
+        mean = IMAGENET_DEFAULT_MEAN
+        std = IMAGENET_DEFAULT_STD
+
+    resize_im = args.input_size > 32
+    if is_train :
+        # this should always dispatch to transforms_imagenet_train
+        transform = create_transform(
+            input_size=args.input_size,
+            is_training=True,
+            color_jitter=args.color_jitter,
+            auto_augment=args.aa,
+            interpolation='bicubic',
+            re_prob=args.reprob,
+            re_mode=args.remode,
+            re_count=args.recount,
+            mean=mean,
+            std=std,
+        )
+        if not resize_im :
+            # replace RandomResizedCropAndInterpolation with
+            # RandomCrop
+            transform.transforms[0] = transforms.RandomCrop(
+                args.input_size, padding=4)
+        return transform
+
+    # eval transform
+    t = []
+    if args.input_size <= 224 :
+        crop_pct = 224 / 256
+    else :
+        crop_pct = 1.0
+    size = int(args.input_size / crop_pct)
+    t.append(
+        transforms.Resize(size, interpolation=PIL.Image.BICUBIC),  # to maintain same ratio w.r.t. 224 images
+    )
+    t.append(transforms.CenterCrop(args.input_size))
+
+    t.append(transforms.ToTensor())
+    t.append(transforms.Normalize(mean, std))
+    return transforms.Compose(t)

src/model_jigsaw.py

@@ -0,0 +1,235 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+import einops
+
+from timm.models.vision_transformer import PatchEmbed, Block
+
+import utils
+
+class MaskedAutoencoderViT(nn.Module):
+    """ Masked Autoencoder with VisionTransformer backbone
+    """
+    def __init__(self,
+                 nb_cls=10,
+                 img_size=224,
+                 patch_size=16,
+                 in_chans=3,
+                 embed_dim=1024,
+                 depth=24,
+                 num_heads=16,
+                 mlp_ratio=4.,
+                 norm_layer=nn.LayerNorm):
+        super().__init__()
+
+        # --------------------------------------------------------------------------
+        # MAE encoder specifics
+        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+        self.num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding
+
+        self.blocks = nn.ModuleList([
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+        self.head = torch.nn.Linear(embed_dim, nb_cls)
+        self.jigsaw = torch.nn.Sequential(*[torch.nn.Linear(embed_dim, embed_dim),
+                                          torch.nn.ReLU(),
+                                          torch.nn.Linear(embed_dim, embed_dim),
+                                          torch.nn.ReLU(),
+                                          torch.nn.Linear(embed_dim, self.num_patches)])
+        self.target = torch.arange(self.num_patches)
+        
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # initialization
+        # initialize (and freeze) pos_embed by sin-cos embedding
+        pos_embed = utils.get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        
+        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.patch_embed.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
+        torch.nn.init.normal_(self.cls_token, std=.02)
+        
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def patchify(self, imgs):
+        """
+        imgs: (N, 3, H, W)
+        x: (N, L, patch_size**2 *3)
+        """
+        p = self.patch_embed.patch_size[0]
+        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
+
+        h = w = imgs.shape[2] // p
+        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
+        x = torch.einsum('nchpwq->nhwpqc', x)
+        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))
+        return x
+
+    def unpatchify(self, x):
+        """
+        x: (N, L, patch_size**2 *3)
+        imgs: (N, 3, H, W)
+        """
+        p = self.patch_embed.patch_size[0]
+        h = w = int(x.shape[1]**.5)
+        assert h * w == x.shape[1]
+        
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
+        return imgs
+
+    def random_masking(self, x, mask_ratio):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+        
+        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+        
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+        # target = einops.repeat(self.target, 'L -> N L', N=N) 
+        # target = target.to(x.device)
+        
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep] # N, len_keep
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+        target_masked = ids_keep
+        
+        return x_masked, target_masked
+
+    def forward_jigsaw(self, x, mask_ratio):
+        # embed patches
+        x = self.patch_embed(x)
+
+        # masking: length -> length * mask_ratio
+        x, target = self.random_masking(x, mask_ratio)
+
+        # append cls token
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        x = self.jigsaw(x[:, 1:])
+        return x.reshape(-1, self.num_patches), target.reshape(-1)
+    
+    def forward_cls(self, x) : 
+        # embed patches
+        x = self.patch_embed(x)
+
+        # add pos embed w/o cls token
+        x = x + self.pos_embed[:, 1:, :]
+
+        # append cls token
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        x = self.head(x[:, 0])
+        return x
+
+    def forward(self, x_jigsaw, x_cls, mask_ratio) :
+        pred_jigsaw, targets_jigsaw = self.forward_jigsaw(x_jigsaw, mask_ratio)
+        pred_cls = self.forward_cls(x_cls)
+        return pred_jigsaw, targets_jigsaw, pred_cls
+        
+
+def mae_vit_small_patch16(nb_cls, **kwargs):
+    model = MaskedAutoencoderViT(nb_cls,
+                                 img_size=224,
+                                 patch_size=16,
+                                 embed_dim=384,
+                                 depth=12,
+                                 num_heads=6,
+                                 mlp_ratio=4,
+                                 norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                                 **kwargs)
+    return model
+
+
+def mae_vit_base_patch16(nb_cls, **kwargs):
+    model = MaskedAutoencoderViT(nb_cls,
+                                 img_size=224,
+                                 patch_size=16,
+                                 embed_dim=768,
+                                 depth=12,
+                                 num_heads=12,
+                                 mlp_ratio=4,
+                                 norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                                 **kwargs)
+    return model
+
+
+def mae_vit_large_patch16(nb_cls, **kwargs):
+    model = MaskedAutoencoderViT(nb_cls,
+                                 img_size=224,
+                                 patch_size=16,
+                                 embed_dim=1024,
+                                 depth=24,
+                                 num_heads=16,
+                                 mlp_ratio=4,
+                                 norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                                 **kwargs)
+    return model
+
+
+def create_model(arch, nb_cls) :     
+    if arch == 'vit_small_patch16' : 
+        return mae_vit_small_patch16(nb_cls)
+    elif arch == 'vit_base_patch16' : 
+        return mae_vit_base_patch16(nb_cls)
+    elif arch == 'vit_large_patch16' : 
+        return mae_vit_large_patch16(nb_cls)
+    
+if __name__ == '__main__':
+    
+    net = create_model(arch = 'vit_small_patch16', nb_cls = 10) 
+    net = net.cpu()  # Move the model to CPU instead of CUDA
+    img = torch.cuda.FloatTensor(6, 3, 224, 224)
+    mask_ratio = 0.75
+    with torch.no_grad():
+        x, target = net.forward_jigsaw(img, mask_ratio)

src/utils.py

@@ -0,0 +1,267 @@
+'''Some helper functions for PyTorch, including:
+    - get_mean_and_std: calculate the mean and std value of dataset.
+    - msr_init: net parameter initialization.
+    - progress_bar: progress bar mimic xlua.progress.
+'''
+import os
+import sys
+import time
+import math
+
+import logging
+from datetime import datetime
+import torch
+import numpy as np
+from torch.nn import Parameter
+
+
+def get_logger(out_dir):
+    logger = logging.getLogger('Exp')
+    logger.setLevel(logging.INFO)
+    formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
+
+    ts = str(datetime.now()).split(".")[0].replace(" ", "_")
+    ts = ts.replace(":", "_").replace("-", "_")
+    file_path = os.path.join(out_dir, "run_{}.log".format(ts)) if os.path.isdir(out_dir) else out_dir.replace('.pth.tar', '')
+    file_hdlr = logging.FileHandler(file_path)
+    file_hdlr.setFormatter(formatter)
+
+    strm_hdlr = logging.StreamHandler(sys.stdout)
+    strm_hdlr.setFormatter(formatter)
+
+    logger.addHandler(file_hdlr)
+    logger.addHandler(strm_hdlr)
+    return logger
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size()[0]
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
+            res.append(correct_k.mul_(100.0 / batch_size))
+        return res
+
+
+def update_lr(iteration, warmup_iter, total_iter, max_lr, min_lr) :
+    if iteration < warmup_iter:
+        current_lr = max_lr * iteration / warmup_iter 
+    else:
+        current_lr = min_lr + (max_lr - min_lr) * 0.5 * \
+            (1. + math.cos(math.pi * (iteration - warmup_iter) / (total_iter - warmup_iter)))
+    return current_lr
+
+
+def adjust_learning_rate(optimizer, iteration, warmup_iter, total_iter, max_lr, min_lr):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    current_lr = update_lr(iteration, warmup_iter, total_iter, max_lr, min_lr)
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = current_lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = current_lr
+    return current_lr
+
+    
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed
+
+
+def load_my_state_dict(net, state_dict):
+    own_state = net.state_dict()
+    for name, param in state_dict.items():
+        name = name.replace('module.','')
+        if name not in own_state:
+            continue
+        if isinstance(param, Parameter):
+            # backwards compatibility for serialized parameters
+            param = param.data
+        own_state[name].copy_(param)
+
+# Fix for non-interactive environments
+try:
+    _, term_width = os.popen('stty size', 'r').read().split()
+    term_width = int(term_width)
+except ValueError:
+    term_width = 80  # Set a default value if the stty command fails
+
+TOTAL_BAR_LENGTH = 65.
+last_time = time.time()
+begin_time = last_time
+
+
+def progress_bar(current, total, msg=None):
+    global last_time, begin_time
+    if current == 0:
+        begin_time = time.time()  # reset for new bar.
+
+    cur_len = int(TOTAL_BAR_LENGTH*current/total)
+    rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1
+
+    sys.stdout.write(' [')
+    for i in range(cur_len):
+        sys.stdout.write('=')
+    sys.stdout.write('>')
+    for i in range(rest_len):
+        sys.stdout.write('.')
+    sys.stdout.write(']')
+
+    cur_time = time.time()
+    step_time = cur_time - last_time
+    last_time = cur_time
+    tot_time = cur_time - begin_time
+
+    L = []
+    L.append('  Step: %s' % format_time(step_time))
+    L.append(' | Tot: %s' % format_time(tot_time))
+    if msg:
+        L.append(' | ' + msg)
+
+    msg = ''.join(L)
+    sys.stdout.write(msg)
+    for i in range(term_width-int(TOTAL_BAR_LENGTH)-len(msg)-3):
+        sys.stdout.write(' ')
+
+    # Go back to the center of the bar.
+    for i in range(term_width-int(TOTAL_BAR_LENGTH/2)+2):
+        sys.stdout.write('\b')
+    sys.stdout.write(' %d/%d ' % (current+1, total))
+
+    if current < total-1:
+        sys.stdout.write('\r')
+    else:
+        sys.stdout.write('\n')
+    sys.stdout.flush()
+
+
+def format_time(seconds):
+    days = int(seconds / 3600/24)
+    seconds = seconds - days*3600*24
+    hours = int(seconds / 3600)
+    seconds = seconds - hours*3600
+    minutes = int(seconds / 60)
+    seconds = seconds - minutes*60
+    secondsf = int(seconds)
+    seconds = seconds - secondsf
+    millis = int(seconds*1000)
+
+    f = ''
+    i = 1
+    if days > 0:
+        f += str(days) + 'D'
+        i += 1
+    if hours > 0 and i <= 2:
+        f += str(hours) + 'h'
+        i += 1
+    if minutes > 0 and i <= 2:
+        f += str(minutes) + 'm'
+        i += 1
+    if secondsf > 0 and i <= 2:
+        f += str(secondsf) + 's'
+        i += 1
+    if millis > 0 and i <= 2:
+        f += str(millis) + 'ms'
+        i += 1
+    if f == '':
+        f = '0ms'
+    return f