Output Unknown error handling prediction.

[swapnil-lader]

Hi, I am trying to return an image using model prediction and it's throwing me an unknown error, In local I am able to get the output out of the model but in replicate it throws the above error.

##Base Code

from cog import BasePredictor, Input, Path,File
import os
import cv2
from PIL import Image
import pandas as pd
import os
import torch
from src.dataset import DatasetInference
from src.model import ModelAlpha, Model
from models.alpha_model_config import cfg as alpha_model_cfg
from models.trimap_model_config import cfg as trimap_model_cfg
from tqdm import tqdm
import numpy as np
import pandas as pd
import cv2
import glob
import torch
from torch.utils.data import DataLoader
import albumentations as alb
from os import path
import tempfile
os.environ['TRANSFORMERS_CACHE'] = './cache/'
os.environ['TORCH_HOME'] = './cache/'

class Predictor(BasePredictor):

def setup(self):
    self.device = torch.device("cpu")
    self.alpha_model_general , self.trimap_model_general = self.get_models("models/alpha_model_general.ckpt" , "models/trimap_model_general.ckpt" , self.device)

def remove_data(self):
    normal_images = os.listdir("data/output/val/images/")
    mask_path = "output_mask/"
    mask_images = os.listdir(mask_path)

    for mask , normal in zip(mask_images , normal_images):
        os.remove(mask_path + mask)
        os.remove("data/output/val/images/" + normal)

    print("data cleaning complete!")
    
def create_csv(self,image_name , tracking_id):    
    test_dataset = {
        "image_path" : [image_name],
        "id" : [tracking_id]
    }
    dataframe = pd.DataFrame(test_dataset)
    csv_file = "data/test_dataset_" + tracking_id + ".csv"
    dataframe.to_csv(csv_file)
    # print("CSV file created!")
    return csv_file

def get_models(self,alpha_path , trimap_path , device = torch.device("cpu")):
    # print('Load TriMap Model...')
    trimap_model = Model(cfg=trimap_model_cfg)
    trimap_model = trimap_model.load_from_checkpoint(
        trimap_path, cfg=trimap_model_cfg)
    trimap_model.freeze()
    trimap_model = trimap_model.to(device)
    # print('OK!\n')

    # print('Load Alpha Model...')
    alpha_model = ModelAlpha(cfg=alpha_model_cfg)
    alpha_model = alpha_model.load_from_checkpoint(
        alpha_path, cfg=alpha_model_cfg)
    alpha_model.freeze()
    alpha_model = alpha_model.to(device)
    # print('OK!\n')

    return alpha_model , trimap_model

def blend(self, image_path, alpha_path):   
    foreground = cv2.imread(image_path)
    alpha = cv2.imread(alpha_path)
    background = np.ones(foreground.shape)

    foreground = foreground.astype(float)
    background = background.astype(float)

    alpha = alpha.astype(float)/255
    
    # print(foreground.shape , alpha.shape)

    foreground = cv2.multiply(alpha, foreground)
    background = cv2.multiply(1.0 - alpha, background)

    blend = cv2.add(foreground, background)

    return blend

def prepare_dataset(self,**kwargs):
    src_df = pd.read_csv(kwargs['image_csv_path'])
    # print(f'Alpha masks will be built for {len(src_df)} image(s).')

    #  get path to save results
    dst_dir = kwargs['mask_path']
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir, exist_ok=True)

    # print(f'Alpha masks will be saved in {dst_dir}\n')

    output_dir = kwargs["output_path"]
    if not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    # print('Load the models...')

    # print('Create dataset...')
    t = alb.Compose([
        alb.Resize(512, 512),
        alb.Normalize(mean=trimap_model_cfg.MEAN, std=trimap_model_cfg.STD)])

    test_dataset = DatasetInference(df=src_df, transform=t)

    test_dataloader = DataLoader(
        dataset=test_dataset,
        batch_size=kwargs.get('batch_size', 1),
        shuffle=False,
        num_workers=kwargs.get('workers', 2))
    # print('OK!\n')

    #  iterate over data and build alphas
    # print('Get alpha masks...')

    return test_dataloader , dst_dir , output_dir

def predict_output(self,dataset , test_dataloader  , dst_dir, output_dir , device = torch.device("cpu")):
    public_urls = []
    mask_paths = []
    k = 0

    for batch in tqdm(test_dataloader, total=len(test_dataloader)):
        # print(batch)
        img, data = batch
        img = img.to(device)
        # print(img)
        # print(img.shape)
        trimap = self.trimap_model_general(img)
        trimap = self.trimap_model_general.activation(trimap)
        trimap = trimap.argmax(1, keepdim=True)
        trimap_processed = trimap.clone().detach()
        trimap_processed[trimap_processed == 1] = 255
        trimap_processed[trimap_processed == 2] = 128
        trimap_processed = trimap_processed.to(torch.float32)
        trimap_processed = trimap_processed / 255

        img_trimap = torch.cat([img, trimap_processed], dim=1)
        trimap_output = trimap_processed.detach().cpu().numpy()*255
        trimap_output = np.reshape(trimap_output , (trimap_output.shape[0] , trimap_output.shape[2] , trimap_output.shape[3] , trimap_output.shape[1]))

        alpha = self.alpha_model_general(img_trimap)
        alpha = self.alpha_model_general.activation(alpha['refine_output'])

        trimap_final = alpha
        k+=1
        # print("reached for loop")
        for i in range(len(img)):
            # print(dataset['id'][i] , f'_{k}_{i}.png')
            mask_path = os.path.join(
                dst_dir,
                dataset['id'][i] + f'_{k}_{i}.png')

            m = trimap_final[i].cpu().detach().numpy()[0, ...]
            m = np.uint8(m * 255)

            m = cv2.resize(m, (int(data['size'][0][i]), int(data['size'][1][i] )))
            # m = np.expand_dims(m, axis=2)

            numpy_image = img[i].cpu().detach().numpy()
            
            numpy_image = np.reshape(numpy_image , (numpy_image.shape[1] ,numpy_image.shape[2] ,numpy_image.shape[0]))

            resized_image = cv2.resize(numpy_image , (int(data['size'][0][i]), int(data['size'][1][i])))
            # print(resized_image.shape , m.shape)

            # blended_image = blend(resized_image , m)

            cv2.imwrite(f"{mask_path}" , m)
            # cv2.imwrite(mask_path, m)

            mask_paths.append(mask_path)

        image_path = sorted(glob.glob('data/output/val/images/*'))
        alpha_path = sorted(glob.glob(dst_dir + '/*'))


        for i in tqdm(image_path):
            # try:
            ID = dataset[dataset["image_path"]==i]["id"].to_string().split(" ")[-1]
            # print(dataset[dataset["image_path"]==i]["id"].to_string().split(" ")[-1])

            for j in tqdm(alpha_path):
                if ID in j: 
                    blended_image =  self.blend(i, j)
                    output_path = Path(tempfile.mkdtemp()) / "output.png"
                    cv2.imwrite(str(output_path) , blended_image)

                    # os.remove(output_location)

            # except Exception as e:
            #     print(e)

        self.remove_data()

        return output_path

# The arguments and types the model takes as input
def predict(self,
        image: Path = Input(description="Image to run inference on"),
        tracking_id: str = Input(description="Insert Your tracking id",default="test")
) -> Path:
    """Run a single prediction on the model"""
    filepath = str(image)
    tracking_id = tracking_id
    device = "cpu"
    raw_image = Image.open(filepath).convert("RGB")
    input_image_dir = "data/output/val/images"
    if path.exists(input_image_dir):
        pass
    else:
        os.makedirs(input_image_dir)
    image_download_path = input_image_dir+"/"+"test.png"
    # print(image_download_path)
    raw_image.save(image_download_path)
    
    csv_path = self.create_csv(image_download_path , tracking_id)
    # print(csv_path)
    if path.exists("blended_images/") and path.exists("output_mask/"):
        pass
    else:
        os.makedirs("blended_images/")
        os.makedirs("output_mask/")
    req = {}    
    req["image_csv_path"] = csv_path
    req["mask_path"] = "output_mask/"
    req["output_path"] = "output/"
    test_dataloader , dst_dir , output_dir = self.prepare_dataset(**req)   
    dataset = pd.read_csv(req["image_csv_path"])
    
    # print(test_dataloader , dst_dir , output_dir , dataset.head())
    
    output_location = self.predict_output(
        dataset,
        test_dataloader,
        dst_dir,
        output_dir,
        device)
    # print(output_location,type(output_location))
    return output_location