Add application file

app.py

@@ -0,0 +1,332 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+import gradio as gr
+import torch.nn as nn
+import torchvision.transforms as transforms
+
+F = torch.nn.functional
+
+class ConvLSTMCell(nn.Module):
+    def __init__(self, input_channels, hidden_channels, kernel_size, bias=True):
+        super(ConvLSTMCell, self).__init__()
+
+        self.input_channels = input_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+
+        self.conv_ii = nn.Conv2d(
+            self.input_channels,
+            self.hidden_channels,
+            self.kernel_size,
+            padding=self.kernel_size // 2,
+            bias=bias,
+        )
+        self.conv_hi = nn.Conv2d(
+            self.hidden_channels,
+            self.hidden_channels,
+            self.kernel_size,
+            padding=self.kernel_size // 2,
+            bias=bias,
+        )
+
+        self.conv_if = nn.Conv2d(
+            self.input_channels,
+            self.hidden_channels,
+            self.kernel_size,
+            padding=self.kernel_size // 2,
+            bias=bias,
+        )
+        self.conv_hf = nn.Conv2d(
+            self.hidden_channels,
+            self.hidden_channels,
+            self.kernel_size,
+            padding=self.kernel_size // 2,
+            bias=bias,
+        )
+
+        self.conv_ig = nn.Conv2d(
+            self.input_channels,
+            self.hidden_channels,
+            self.kernel_size,
+            padding=self.kernel_size // 2,
+            bias=bias,
+        )
+        self.conv_hg = nn.Conv2d(
+            self.hidden_channels,
+            self.hidden_channels,
+            self.kernel_size,
+            padding=self.kernel_size // 2,
+            bias=bias,
+        )
+
+        self.conv_io = nn.Conv2d(
+            self.input_channels,
+            self.hidden_channels,
+            self.kernel_size,
+            padding=self.kernel_size // 2,
+            bias=bias,
+        )
+        self.conv_ho = nn.Conv2d(
+            self.hidden_channels,
+            self.hidden_channels,
+            self.kernel_size,
+            padding=self.kernel_size // 2,
+            bias=bias,
+        )
+
+    def forward(self, x, hidden_state):
+        h_prev, c_prev = hidden_state
+
+        i = torch.sigmoid(self.conv_ii(x) + self.conv_hi(h_prev))
+        f = torch.sigmoid(self.conv_if(x) + self.conv_hf(h_prev))
+        g = F.relu(self.conv_ig(x) + self.conv_hg(h_prev))
+        o = torch.sigmoid(self.conv_io(x) + self.conv_ho(h_prev))
+        c = f * c_prev + i * g
+        h = o * F.relu(c)
+
+        return h, c
+
+class ConvLSTM(nn.Module):
+    def __init__(self, input_channels, hidden_channels, kernel_size, bias=True):
+        super(ConvLSTM, self).__init__()
+
+        self.input_channels = input_channels
+        self.hidden_channels = hidden_channels
+
+        # Single ConvLSTM layer
+        self.conv_lstm_cell = ConvLSTMCell(
+            self.input_channels, self.hidden_channels, kernel_size, bias
+        )
+
+    def forward(self, x):
+        batch_size, channels, sequence_length, height, width = x.size()
+
+        # Initialize hidden state and cell state
+        h = torch.zeros(batch_size, self.hidden_channels, height, width).to(x.device)
+        c = torch.zeros(batch_size, self.hidden_channels, height, width).to(x.device)
+
+        outputs = list()
+
+        # Process each time step in the sequence
+        for t in range(sequence_length):
+            h, c = self.conv_lstm_cell(x[:, :, t, :, :], (h, c))
+            outputs.append(h)
+
+        outputs = torch.stack(outputs, dim=0).permute(1, 2, 0, 3, 4).contiguous()
+
+        return outputs
+
+class NextFramePredictionModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        val = 256
+        self.convlstm0 = nn.Sequential(
+            ConvLSTM(3, val, 5),  # Modified line
+            nn.BatchNorm3d(val),
+        )
+        self.convlstm1 = nn.Sequential(
+            ConvLSTM(val, val, 3),
+            nn.BatchNorm3d(val),
+        )
+        self.convlstm2 = nn.Sequential(
+            ConvLSTM(val, val, 1),
+            nn.BatchNorm3d(val),
+        )
+        self.final = ConvLSTM(val, 3, 1)
+
+    def forward(self, x):
+        x = self.convlstm0(x)
+        x = self.convlstm1(x)
+        x = self.convlstm2(x)
+
+        return self.final(x)
+
+class ModelWrapper(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.arch = NextFramePredictionModel()
+
+
+    def forward(self, x):
+        return self.arch(x)
+
+def preprocess_image(image):
+    """
+    Preprocesses the input image to be compatible with the model.
+
+    Args:
+        image_path (str): Path to the input image.
+
+    Returns:
+        torch.Tensor: Preprocessed image tensor.
+    """
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),  # Default size as per training
+        transforms.ToTensor(),
+        #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    image = Image.fromarray(image).convert('RGB')
+    image = transform(image)
+    image = image.unsqueeze(0)  # Add batch dimension
+    image = image.permute(1, 0, 2, 3)  # Rearrange dimensions as per training setup
+    return image.unsqueeze(0)
+
+
+def preprocess_image_no_normalize(image_path: str):
+    """
+    Preprocesses the input image to be compatible with the model.
+
+    Args:
+        image_path (str): Path to the input image.
+
+    Returns:
+        torch.Tensor: Preprocessed image tensor.
+    """
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),  # Default size as per training
+        transforms.ToTensor(),
+    ])
+
+    image = Image.open(image_path).convert("RGB")
+    image = transform(image)
+    image = image.unsqueeze(0)  # Add batch dimension
+    image = image.permute(1, 0, 2, 3)  # Rearrange dimensions as per training setup
+    return image.unsqueeze(0)
+
+
+def denormalize_image(output_image: torch.Tensor):
+    """
+    Denormalizes the output image from model predictions.
+
+    Args:
+        output_image (torch.Tensor): The model's raw output image tensor in shape (H, W, C).
+
+    Returns:
+        torch.Tensor: The denormalized image tensor in shape (H, W, C).
+    """
+    # Check if the input image is in HWC format and convert to CHW format
+    if output_image.ndimension() == 3 and output_image.shape[2] == 3:
+        output_image = output_image.permute(2, 0, 1)  # Convert to C x H x W format
+
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)  # Shape (3, 1, 1)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)  # Shape (3, 1, 1)
+
+    # Reverse normalization: output_image * std + mean
+    denormalized_image = output_image * std + mean
+
+    # Convert back to HWC format for visualization
+    denormalized_image = denormalized_image.permute(1, 2, 0)  # Convert back to H x W x C
+    return denormalized_image
+
+
+
+
+
+def load_model(model_path: str, device: str):
+    """
+    Load the trained NextFramePredictionModel from the specified path.
+
+    Args:
+        model_path (str): Path to the saved model file (e.g., mode.pth).
+        device (str): Device to load the model on (e.g., 'cpu' or 'cuda').
+
+    Returns:
+        torch.nn.Module: The loaded model in evaluation mode.
+    """
+
+
+    # Initialize the model
+    model = ModelWrapper()
+    model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
+    model.eval()  # Set the model to evaluation mode
+    return model
+
+
+def calculate_time_steps(temperature: float, base_temperature: float = 25, Q10: float = 2):
+    """
+    Calculates the equivalent time steps needed based on the given temperature.
+
+    Args:
+        temperature (float): The current temperature.
+        base_temperature (float): The temperature for which the model is calibrated (default is 25).
+        Q10 (float): The Q10 coefficient (default is 2).
+
+    Returns:
+        int: The number of prediction steps needed.
+    """
+    k1 = 1  # Original spoilage rate at base_temperature (1 step per day at 25°C)
+    k2 = k1 * Q10 ** ((temperature - base_temperature) / 10)
+    return max(1, round(k2))  # Ensure at least 1 step
+
+
+def predict_next_frame(image, model: torch.nn.Module, num_steps: int = 1):
+    """
+    Predicts the next frame(s) based on the input image and temperature-adjusted steps.
+
+    Args:
+        image_path (str): Path to the input image.
+        model (torch.nn.Module): Loaded PyTorch model.
+        num_steps (int): Number of prediction steps to perform.
+
+    Returns:
+        np.ndarray: Predicted frame as a NumPy array after `num_steps` iterations.
+    """
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.eval()
+
+    # Preprocess the input image
+    input_tensor = preprocess_image(image).to(device)
+
+    # Iteratively predict the next frame
+    for _ in range(num_steps):
+        with torch.no_grad():
+            output_tensor = model(input_tensor)
+        # Update input_tensor for the next prediction
+        if _ == num_steps-1:
+            output_frame = output_tensor.permute(0, 2, 3, 4, 1)[0][0].detach().cpu().numpy()
+        input_tensor = output_tensor
+
+    # Postprocess the final output
+    return torch.tensor(output_frame) #denormalize_image(torch.tensor(output_frame))
+
+
+def load_and_predict(image, temperature: float=25, model_path: str = 'model (2).pth'):
+    """
+    Loads the model, calculates time steps, and predicts the next frame for the given image and temperature.
+
+    Args:
+        image_path (str): Path to the input image.
+        temperature (float): The current temperature.
+        model_path (str): Path to the saved model file.
+
+    Returns:
+        np.ndarray: Predicted frame as a NumPy array.
+    """
+    # Determine the device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Load the model
+    model = load_model(model_path, device)
+
+    # Calculate the number of steps based on temperature
+    num_steps = calculate_time_steps(temperature)
+    print(num_steps)
+    # Predict the next frame(s)
+    return predict_next_frame(image, model, num_steps=num_steps)
+
+
+
+# Gradio interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy"),
+    outputs="text",
+    title="Banana Predictor",
+    description="Upload image.",
+)
+
+if __name__ == "__main__":
+    interface.launch()

model (2).pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3d85e1e5bf3b1f05645c96233dea360c07880de846cb580ed5874121a66f7aa
+size 47567427

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+torch
+torchvision
+Pillow