main.py

# --------------------------------------------------------
# ALPS
# Copyright (c)
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
import argparse
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
import sys
import os
import glob
import json
import gc

sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

import torchvision
from torchvision import models
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
from torchvision import transforms as T

from sklearn import datasets
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import MiniBatchKMeans
from sklearn.decomposition import PCA
import joblib

from sam_vit import ImageEncoderViT
from functools import partial

import time
import datetime

import os
import os.path as osp

from utils.logger import create_logger
from utils.utils import save_anns, filter_masks, write_masks_to_folder, create_color_mapping, add_mask_to_template, add_gt_seg_mask_to_template
from utils.model import build_sam_vit
from utils.dataset import Custom_Dataset, Mask_Vec_Dataset, Mask_Vec_TestDataset

# GPU Selection
os.environ['CUDA_VISIBLE_DEVICES']="0"


# Strict Requirement！！！
"""
Dataset Structure 
├── img_dir
│   ├── train
│   │   ├── xxx{img_suffix}
│   │   ├── yyy{img_suffix}
│   │   ├── zzz{img_suffix}
│   │   ├── ....
│   ├── val
│   │   ├── xxx{img_suffix}
│   │   ├── yyy{img_suffix}
│   │   ├── zzz{img_suffix}
│   │   ├── ....
│   ├── test
│   │   ├── xxx{img_suffix}
│   │   ├── yyy{img_suffix}
│   │   ├── zzz{img_suffix}
│   │   ├── ....
"""

# ---------Parameter Configuration---------
parser = argparse.ArgumentParser('ALPS', add_help=False)
## Modification Necessarily
### Normal Configuration
parser.add_argument('--root_dir', default="",type=str, metavar="FILE", help='path to the dataset dir', )
parser.add_argument('--process_list', type=list, default=['train', 'val'], choices=['train', 'val', 'test'])
parser.add_argument('--img_name', type=str, default='img_dir', help='img name')
parser.add_argument('--threshold', type=float, default=0.3, help='the threshold value to filter the useless masks')
parser.add_argument('--label_dir', type=str, default="instance_mask", help='path to the final pseudo masks obtained by SAM and K-Means')
parser.add_argument('--logger_dir', type=str, default="logs", help='path to log dir')
parser.add_argument('--logger_name', type=str, default="Test", help='log name')
parser.add_argument('--temp_dir', type=str, default='temporary_files', help='path to save the temporary files')
parser.add_argument('--image_suffix', type=str, default='.png', help='the suffix of data source, e.g., .png \ .jpg ...')
parser.add_argument('--delete_temp', type=bool, default=False, choices=[True, False], help='if True the program will delete the tempoary files generated by this script !')

### SAM Configuration
parser.add_argument('--sam_checkpoint', default="",type=str, help='checkpoint path to the sam pre-trained model', )
parser.add_argument('--model_type', type=str, default='vit_h', choices=['vit_b', 'vit_l', 'vit_h'])
parser.add_argument('--device', type=str, default='cuda', choices=['cuda', 'cpu'])
### Online K-Means
parser.add_argument('--number_clusters', type=int, default=46, help='the number of clustering classes')
parser.add_argument('--batch_size', type=int, default=512, help='the batch size adopted in Online K-means')
parser.add_argument('--kmeans_save_name', type=str, default='k_means.pkl', help='model save name of online k-means')
### Visualization
parser.add_argument('--vis', type=bool, default=False, choices=[True, False], help='if True the program will meanwhile save the visualization pseudo masks, which will occupy more storages!')
parser.add_argument('--vis_dir', type=str, default="vis_mask", help='path to the visualization instance masks')
parser.add_argument('--vis_prefix', type=str, default='instance_vis', help='the suffix of saving pseudo label masks')

## Modification Unnecessarily!!!  These are temporary files.
parser.add_argument('--mask_save_dir_name', type=str, default='test_mask', help='the temporary file to save the masks generated by SAM')
parser.add_argument('--mask_save_dir_prefix', type=str, default='all_mask', help='the prefix name of each mask saving dir')
parser.add_argument('--json_save_dir_name', type=str, default='mask_vec', help='the feature vector of masks to save as a json file')
parser.add_argument('--npy_save_dir_name', type=str, default='mask_vec_npy', help='the feature vector of masks to save as a npy file')
parser.add_argument('--instance_label_name', type=str, default='mask2instance', help='the dict of mask to instance label')
parser.add_argument('--instance_save_dir_name', type=str, default='instance_dir', help='the path to save the dict')
parser.add_argument('--instance_path', type=str, default='instance_mask_visualization', help='visualizaiton instance label name')
parser.add_argument('--instance_prefix', type=str, default='instance_color', help='the suffix of saving pseudo label masks')

args = parser.parse_args()

# Logger
os.makedirs(osp.join(args.root_dir, args.logger_dir), exist_ok=True)
logger = create_logger(output_dir=osp.join(args.root_dir, args.logger_dir), name=args.logger_name)

# Step 1
def generate_all_masks():
    mask_save_dir_name = osp.join(args.root_dir, args.temp_dir, args.mask_save_dir_name)
    dataset_path = osp.join(args.root_dir, args.img_name)
    # mk dir
    os.makedirs(mask_save_dir_name, exist_ok=True)
    for item in args.process_list:
        save_path_dir = os.path.join(mask_save_dir_name, item)
        os.makedirs(save_path_dir, exist_ok=True)
        logger.info('Saving dir create finished!')
    
    # create sam
    sam = sam_model_registry[args.model_type](checkpoint=args.sam_checkpoint)
    sam.to(device=args.device)
    mask_generator = SamAutomaticMaskGenerator(sam)
    logger.info('SAM Model Loaded!')
    
    # load image name
    for item in args.process_list:
        path = os.path.join(dataset_path, item)
        img_list = glob.glob(os.path.join(path, '*'+args.image_suffix))
        for img in img_list:
            # P0000_0_896_0_896
            img_name = img.split('/')[-1].split('.')[0] # obtain img name
            # mk masks save dir for each img
            # logger.info(img_name)
            all_mask_dir_name = img_name + '_' + args.mask_save_dir_prefix
            all_mask_dir_path = os.path.join(mask_save_dir_name, item, all_mask_dir_name)
            os.makedirs(all_mask_dir_path, exist_ok=True)
            # load image
            img_path = os.path.join(dataset_path, img)
            image = cv2.imread(img_path)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            # Generate Mask
            masks = mask_generator.generate(image)
            # mask save path
            write_masks_to_folder(image, masks, all_mask_dir_path, args.image_suffix, args.threshold)
            logger.info('Save all masks of '+ str(img_name) +' as ' + all_mask_dir_path + ' successed !')
            del image
            del masks
            gc.collect()
    logger.info('Job finished!')


# Step 2
def use_encoder_align_mask_feature():
    logger.info(f'Creating the {args.model_type} ...')
    model = build_sam_vit(args.model_type)
    with open(args.sam_checkpoint, "rb") as f:
        state_dict = torch.load(f)
        if any([True if 'image_encoder.' in k else False for k in state_dict.keys()]):
            state_dict = {k.replace('image_encoder.', ''): v for k, v in state_dict.items() if k.startswith('image_encoder.')}
    model.load_state_dict(state_dict)
    logger.info(f'Loading ViT encoder~({args.model_type}) of SAM From {args.sam_checkpoint}...')
    # put it on cuda
    # if torch.cuda.is_available():
    #     model = model.cuda()
    model.to(device=args.device)
    model.eval()

    dataset_path = osp.join(args.root_dir, args.img_name)
    mask_save_dir_name = osp.join(args.root_dir, args.temp_dir, args.mask_save_dir_name)
    
    for item in args.process_list:
        # dataset
        test_dataset = Custom_Dataset(root_dir=os.path.join(dataset_path, item), mask_dir=os.path.join(mask_save_dir_name, item), dataset_path=dataset_path, prefix=args.mask_save_dir_prefix, image_suffix=args.image_suffix)
        test_dataloader = DataLoader(dataset=test_dataset, batch_size=1)
        # json dict save path
        json_mask_avg_vec_save_path = os.path.join(args.root_dir, args.temp_dir, args.json_save_dir_name, item)
        os.makedirs(json_mask_avg_vec_save_path, exist_ok=True)
        # npy save dir
        mask_vec_npy_dir = os.path.join(args.root_dir, args.temp_dir, args.npy_save_dir_name, item)
        os.makedirs(mask_vec_npy_dir, exist_ok=True)

        # inference
        with torch.no_grad():
            for idx, item in enumerate(test_dataloader):
                img_info = dict() # the image information saving dict
                image_name = item['image_name'][0]
                image = item['image']
                masks = item['masks']
                if args.device == 'cuda':
                    image = image.cuda()
                # image.to(device=args.device)
                # print(image.shape)
                features = model(image) # [B, C, H, W]
                # print(features.shape)
                b, c, h, w = features.shape

                masked_feature_list = dict()
                for item in masks:
                    mask_name = item['mask_name'][0]
                    mask = item['mask']
                    # print(mask[mask!=0])
                    if args.device == 'cuda':
                        mask = mask.cuda()
                    # mask.to(device=args.device)

                    # upsampling features and downsample masks
                    # print(mask.shape)
                    rescale_factor = 4 # 4x upsample / down sample
                    t1, t2 = int(mask.shape[2] / rescale_factor), int(mask.shape[3] / rescale_factor)
                    features_rescale = F.interpolate(features, size=[t1, t2], mode='bilinear')
                    mask_rescale = F.interpolate(mask, size=[t1, t2], mode='bilinear')
                    # print('mask_rescale:'+str(torch.count_nonzero(mask_rescale.view(1, 1, -1), dim=2).sum()))
                    masked_feature = torch.mul(features_rescale, mask_rescale)
                    masked_feature = masked_feature.view(b, c, -1)
                    non_zero_count = torch.count_nonzero(masked_feature, dim=2)
                    # calculate h×w avg vec
                    masked_avg_vec = masked_feature.sum(dim=2) / non_zero_count # 加和之后求均值
                    # 将tensor中的nan替换为0
                    masked_avg_vec[torch.isnan(masked_avg_vec)] = 0
                    npy_data = masked_avg_vec.detach().cpu().numpy()
                    single_mask_vec_path = os.path.join(mask_vec_npy_dir, image_name+'_'+mask_name+'.npy')
                    
                    np.save(single_mask_vec_path, npy_data)
                    masked_feature_list[mask_name] = single_mask_vec_path

                    del npy_data
                    del masked_avg_vec
                    del features_rescale
                    del masked_feature
                
                img_info['image_name'] = image_name
                img_info['mask_avg_vec'] = masked_feature_list # all masks for each image saving dict

                # save json dict
                json_save_name = os.path.join(json_mask_avg_vec_save_path, image_name+'_mask_vec.json')
                with open(json_save_name, "w") as file:
                    json.dump(img_info, file)

                logger.info(f'Saving image: {image_name} all masks as {json_save_name}.')
        del test_dataloader
        del test_dataset
        gc.collect()

# Step 3
# Training mini batch(online) K-meas for self-constructed dataset
"""
    rm -rf k_means.pkl
"""
def online_k_means_clustering_train():
    # create K-means
    kmeans = MiniBatchKMeans(n_clusters=args.number_clusters, random_state=0, batch_size=args.batch_size, n_init="auto")

    for item in args.process_list:
        # load all mask vecs
        # json_mask_avg_vec_save_path = os.path.join(root_dir, json_save_dir_name, item)
        # all_mask_vec_path = os.path.join(json_mask_avg_vec_save_path, 'all_mask_vec.npy')
        train_dataset = Mask_Vec_Dataset(root_dir=osp.join(args.root_dir, args.temp_dir), json_dir=args.json_save_dir_name, item_dir=item)
        train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, drop_last=False)
        
        for idx, item in enumerate(train_dataloader):
            item = np.squeeze(item, 1) # [B, 2048]
            kmeans = kmeans.partial_fit(item)

    save_path = os.path.join(args.root_dir, args.kmeans_save_name)
    joblib.dump(kmeans, save_path)
    logger.info(f'Saving pretrained k-means model as {save_path}...')
    # logger.info(kmeans.cluster_centers_)

# Step 4
# Prediction + save instance label sets to json for each mask by online k-means model
def online_k_means_clustering_test():
    save_path = os.path.join(args.root_dir, args.kmeans_save_name)
    kmeans = joblib.load(save_path)
    logger.info(f'Loaded pretrained k-means model from {save_path}...')
    logger.info('kmeans cluster centers:' + str(kmeans.cluster_centers_))
    
    for item in args.process_list:
        # json dict save path
        instance_mask_avg_vec_save_path = os.path.join(args.root_dir, args.temp_dir, args.instance_save_dir_name, item)
        os.makedirs(instance_mask_avg_vec_save_path, exist_ok=True)
        # dataset
        test_dataset = Mask_Vec_TestDataset(root_dir=osp.join(args.root_dir, args.temp_dir), json_dir=args.json_save_dir_name, item_dir=item)
        test_dataloader = DataLoader(dataset=test_dataset, batch_size=args.batch_size, drop_last=False)
        # print(len(test_dataset))
        # test
        image_instance_dict = dict()
        for idx, item in enumerate(test_dataloader):
            # print(item["mask_vec"].shape)
            # print(type(item["image_name"])) # list
            mask_vec = np.squeeze(item["mask_vec"], 1)
            image_name = item["image_name"]
            mask_name = item["mask_name"]
            predicts = kmeans.predict(mask_vec) # list
            for idx, key in enumerate(image_name):
                mask_id = mask_name[idx]
                pred = predicts[idx]
                if key in image_instance_dict.keys():
                    image_instance_dict[key][mask_id] = int(pred)
                else:
                    image_instance_dict[key] = dict()
                    image_instance_dict[key][mask_id] = int(pred)

        json_save_name = os.path.join(instance_mask_avg_vec_save_path, args.instance_label_name +'.json')
        with open(json_save_name, "w") as file:
            json.dump(image_instance_dict, file)

        logger.info(f'Saving mask instance label sets as {json_save_name}...')

# Step 5
# Generate instance mask for training based on above obtained temporary data
def generate_gt_instance_mask():
    # If need Visualization, obtaining the color map
    if args.vis:
        color_map = create_color_mapping()
        instance_seg_vis_path = args.vis_dir
    
    instance_seg_gt_path = args.label_dir
    mask_save_dir_name = osp.join(args.root_dir, args.temp_dir, args.mask_save_dir_name)

    for item in args.process_list:
        # create instance label mask dir
        instance_mask_save_path = os.path.join(args.root_dir, instance_seg_gt_path, item)
        os.makedirs(instance_mask_save_path, exist_ok=True)
        # create visualization save dir
        if args.vis:
            visual_mask_save_path = os.path.join(args.root_dir, instance_seg_vis_path, item)
            os.makedirs(visual_mask_save_path, exist_ok=True)
        # load mask2instance.json
        instance_mask_avg_vec_save_path = os.path.join(args.root_dir, args.temp_dir, args.instance_save_dir_name, item)
        json_path = os.path.join(instance_mask_avg_vec_save_path, args.instance_label_name+'.json')
        with open(json_path,'r', encoding='UTF-8') as f:
            data = json.load(f)
        # create instance map
        for image_name in data.keys():
            mask2instance_dict = data[image_name]
            mask_dir = os.path.join(mask_save_dir_name, item, image_name+'_'+args.mask_save_dir_prefix)
            template_save_path = os.path.join(instance_mask_save_path, image_name+'_'+args.instance_prefix+args.image_suffix)
            # skip already generated files
            if os.path.exists(template_save_path):
                print(f"{template_save_path} is existed!")
                continue
            template = None
            vis_template = None
            idx = 0
            for mask_name in mask2instance_dict.keys():
                instance_label = mask2instance_dict[mask_name]
                # obtain color value
                if args.vis:
                    color_value = color_map[instance_label]
                mask_path = os.path.join(mask_dir, mask_name+args.image_suffix)
                mask = np.array(cv2.imread(mask_path, 0)) / 255.0
                if idx == 0:
                    # generate specific-visualization template
                    if args.vis:
                        vis_template = np.zeros((mask.shape[0], mask.shape[1], 3))
                    # generate normal template for training
                    template = np.zeros((mask.shape[0], mask.shape[1], 1)) # shape要为2D！
                    # print(template.shape)
                    idx += 1
                # 叠加mask
                template = add_gt_seg_mask_to_template(mask, template, int(instance_label))
                if args.vis:
                    vis_template = add_mask_to_template(mask, vis_template, color_value)

            # save template
            template_save_path = os.path.join(instance_mask_save_path, image_name+'_'+args.instance_prefix+args.image_suffix)
            cv2.imwrite(template_save_path, template)
            logger.info(f'Saving instance mask as {template_save_path}...')

            # save visual template
            if args.vis:
                vis_template_save_path = os.path.join(visual_mask_save_path, image_name+'_'+args.vis_prefix+args.image_suffix)
                cv2.imwrite(vis_template_save_path, vis_template)
                logger.info(f'Saving visual mask as {vis_template_save_path}...')

# main function
def main():
    logger.info(f"Time: {datetime.datetime.now()}, Script Staring ...")
    
    # Step 1:
    logger.info(f"Time: {datetime.datetime.now()}, Step 1: predict the mask of each image by SAM")
    generate_all_masks()

    # Step 2:
    logger.info(f"Time: {datetime.datetime.now()}, Step 2: use encoder of SAM to extract feature of the img and align the feature for each mask")
    use_encoder_align_mask_feature()

    # Step 3:
    logger.info(f"Time: {datetime.datetime.now()}, Step 3: Training online K-means based on the feature vectors of each mask")
    online_k_means_clustering_train()

    # Step 4:
    logger.info(f"Time: {datetime.datetime.now()}, Step 4: Predicting the class of each mask by pre-trained online K-means model")
    online_k_means_clustering_test()

    # Step 5:
    logger.info(f"Time: {datetime.datetime.now()}, Step 5: Generating the instance mask based on above temporary data")
    generate_gt_instance_mask()

    # # Step 6:
    if args.delete_temp:
        logger.info(f"Time: {datetime.datetime.now()}, Step 6: Deleting the tempoary files ......")
        cmd = 'rm -rf '+ osp.join(args.root_dir, args.temp_dir)
        os.system(cmd)
        logger.info(f"Time: {datetime.datetime.now()}, Delete Finished!!!")

main()