create_submission.py

import sys
import os
import pandas as pd
import numpy as np
import torch
from tqdm import tqdm
from torch.utils.data import DataLoader
from lib.data_generators import Cells_dataset
from lib.utils import run_competition_test, preprocess_inference, get_cells_rle, save_images_batch
from skimage.transform import resize

####################
# Script arguments #
####################

if len(sys.argv) != 3:
    print(f"Usage: python {sys.argv[0]} <PATH_TO_DATA_FOLDER> <PATH_TO_TRAINED_MODEL>")
    sys.exit()
else:
    data_path  = sys.argv[1]
    model_path = sys.argv[2]

##########################
# Check computing device #
##########################

if torch.cuda.is_available():
    n_gpus = torch.cuda.device_count()
    if n_gpus > 1:
        print(f"\n{n_gpus} GPU's available:")
        for gpu_idx in range(n_gpus):
            print(f"\t-At device cuda:{gpu_idx} -> device model = {torch.cuda.get_device_name(gpu_idx)}")
    else:
        print(f"\nCuda available with device {device} -> device model = {torch.cuda.get_device_name(device_slot)}")
    
    # Select a GPU
    device_slot = torch.cuda.current_device()
    device = torch.device(f"cuda:{device_slot}")
else:
    n_gpus = 0
    device = torch.device("cpu")
    print(f"\nCuda is not available, using {device} instead")

#############################
# Load the pretrained model #
#############################

model = torch.load(model_path)
model = model.to(device)

#################
# Run inference #
#################

# Create the dataframe of the submission file
df = pd.DataFrame(columns=["ImageId", "EncodedPixels"])
# To check if there is a case with no mask generated by the net
empty_mask = True

for root, dirs, files in tqdm(os.walk(data_path)):
    for f in files:
        if f.endswith(".png"):
            empty_mask = True
            image_path = os.path.join(root, f)
            # Read the image and create the input tensor
            image_tensor, orig_shape = preprocess_inference(image_path)
            # Get the model prediction
            pred_mask = model(image_tensor.to(device))
            # Resize the predicted mask to the target shape
            resized_pred_mask = resize(np.squeeze(pred_mask.cpu().detach().numpy()), (orig_shape[1], orig_shape[0])) 
            # Extract the rle encoding for each cell submask
            cells_rle_masks = get_cells_rle(resized_pred_mask)
            # For each cell rle create a row in the submission file
            for cell_rle in cells_rle_masks:
                empty_mask = False
                df = df.append({"ImageId": f[:-4], "EncodedPixels": cell_rle}, ignore_index=True)
            if empty_mask:
                print(f"No mask generated for file {image_path}")
                df = df.append({"ImageId": f[:-4], "EncodedPixels": ""}, ignore_index=True)
                

# Save the submission file
df.to_csv(os.path.join(data_path, "my_submission.csv"), index=False)