Using custom trained model for predictions

munikumar: <pre><code class="lang-auto">class LoadImages:
    def __init__(self, prediction_images_path):
        self.prediction_images_path = prediction_images_path

    def load_images_as_dataframe(self):
        prediction_images = glob.glob(os.path.join(self.prediction_images_path, "*.jpg"))
        prediction_images_df = pd.DataFrame(prediction_images, columns=["ImageName"])
        return prediction_images_df


class PredictionDataset(Dataset):
    """
    Convert images and labels to pytorch dataset
    """
    def __init__(self, data_to_dataset):
        self.data = data_to_dataset
        self.mean = [0.5, 0.5, 0.5]
        self.std = [0.5, 0.5, 0.5]
        self.simple_transformation = transforms.Compose([transforms.Resize((1056, 1056),
                                                                           interpolation=transforms.InterpolationMode.BICUBIC),
                                                         transforms.ToTensor(),
                                                         transforms.Normalize(mean=self.mean, std=self.std)])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img = Image.open(self.data.loc[idx, "ImageName"])
        image = self.simple_transformation(img)
        return image


class ModelPredict:
    def __init__(self, prediction_images_path, model, best_weights_path, batch_size=8):
        super().__init__()
        self.prediction_images_path = prediction_images_path
        self.batch_size = batch_size
        self.current_model = model
        self.best_weights_path = best_weights_path

    def check_paths(self):
        if not os.path.exists(self.prediction_images_path):
            raise FileNotFoundError(f"{self.prediction_images_path} doesn't exist")

        if not os.path.exists(self.best_weights_path):
            raise FileNotFoundError(f"{self.best_weights_path} doesn't exist")

    def prepare_data(self):
        prediction_data_df = LoadImages(self.prediction_images_path).load_images_as_dataframe()
        prediction_dataset = PredictionDataset(prediction_data_df)
        prediction_loader = DataLoader(prediction_dataset, batch_size=self.batch_size, shuffle=False, num_workers=2)
        return prediction_data_df, prediction_loader

    def predict_images(self):
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        prediction_data_df, prediction_loader = self.prepare_data()
        prediction_results = []
        current_model = self.current_model
        current_model.load_state_dict(torch.load(self.best_weights_path, device))
        current_model.eval()

        for images in prediction_loader:
            images = images.to(device)
            with torch.no_grad():
                outputs = current_model(images)
                _, predictions = torch.max(outputs, dim=1)
                predictions_output = predictions.cpu().detach().numpy()
            for prediction in predictions_output:
                prediction_results.append(prediction)

        prediction_data_df["Predictions"] = prediction_results
        predictions_dict = dict(zip(prediction_data_df['ImageName'], prediction_data_df["Predictions"]))

        return predictions_dict

    def classify_images_and_save_predictions(self):
        self.check_paths()
        predictions_dict = self.predict_images()
        with open(os.path.join(self.prediction_images_path, "classifications.json"), "w") as file:
            json.dump(predictions_dict, file, indent=1)
</code></pre>
<p>If I used the above code, prediction_results are all the same for the images I am predicting. I have trained the model for multi-class classification with class labels - 0, 1, 2 but the output predictions of the model is 746 which is not relevant at all.</p>
<p>I have tested with 3 images and I have got the prediction_results as [746, 746, 746]. I am unable to figure out why am I getting these outputs rather than in [0, 1, 2]. May I know what’s wrong with the code?</p>