trjgru.py

import tensorflow as tf
from tensorflow.keras import layers, models # type: ignore
import numpy as np


class TrajectoryGRU2D(layers.Layer):
    def __init__(self, filters, kernel_size, return_sequences=True, **kwargs):
        super().__init__(**kwargs)
        self.filters = filters
        self.kernel_size = kernel_size
        self.return_sequences = return_sequences

        # Projection layer to match GRU feature space
        self.input_projection = layers.Conv2D(filters, (1, 1), padding="same")

        # GRU Gates
        self.conv_z = layers.Conv2D(filters, kernel_size, padding="same", activation="sigmoid")
        self.conv_r = layers.Conv2D(filters, kernel_size, padding="same", activation="sigmoid")
        self.conv_h = layers.Conv2D(filters, kernel_size, padding="same", activation="tanh")

        # Motion-based trajectory update
        self.motion_conv = layers.Conv2D(filters, kernel_size, padding="same", activation="tanh")

    def build(self, input_shape):
        # Ensures input_projection is built with the correct input shape
        self.input_projection.build(input_shape[1:])  # Ignore batch dimension
        super().build(input_shape)

    def call(self, inputs):
        # inputs shape: (batch_size, time_steps, height, width, channels)
        batch_size, time_steps, height, width, channels = tf.unstack(tf.shape(inputs))
        time_steps = inputs.shape[1]

        # Initialize hidden state
        h_t = tf.zeros((batch_size, height, width, self.filters))

        # List to store outputs at each time step
        outputs = []

        # Iterate over time steps
        for t in range(time_steps):
            # Get the input at time step t
            x_t = inputs[:, t, :, :, :]

            # Project input to match GRU feature dimension
            x_projected = self.input_projection(x_t)

            # Compute motion-based trajectory update
            motion_update = self.motion_conv(x_projected)

            # Concatenate projected input, previous hidden state, and motion update
            combined = tf.concat([x_projected, h_t, motion_update], axis=-1)

            # Compute GRU gates
            z = self.conv_z(combined)  # Update gate
            r = self.conv_r(combined)  # Reset gate

            # Compute candidate hidden state
            h_tilde = self.conv_h(tf.concat([x_projected, r * h_t], axis=-1))

            # Update hidden state with motion-based trajectory
            h_t = (1 - z) * h_t + z * h_tilde + motion_update  # Add motion update

            # Store the output if return_sequences is True
            if self.return_sequences:
                outputs.append(h_t)

        # Stack outputs along the time dimension if return_sequences is True
        if self.return_sequences:
            outputs = tf.stack(outputs, axis=1)
        else:
            outputs = h_t

        return outputs

    def compute_output_shape(self, input_shape):
        if self.return_sequences:
            return (input_shape[0], input_shape[1], input_shape[2], input_shape[3], self.filters)
        else:
            return (input_shape[0], input_shape[2], input_shape[3], self.filters)


def build_tgru_model(input_shape=(8, 95, 95, 2)):  # (time_steps, height, width, channels)
    input_tensor = layers.Input(shape=input_shape)

    # Apply TGRU Layers
    x = TrajectoryGRU2D(filters=32, kernel_size=(3, 3), return_sequences=True)(input_tensor)
    x = layers.Conv3D(filters=32, kernel_size=(3, 3, 3), padding='same', activation='relu')(x)
    x = layers.MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding='same')(x)

    x = TrajectoryGRU2D(filters=64, kernel_size=(3, 3), return_sequences=True)(x)
    x = layers.Conv3D(filters=64, kernel_size=(3, 3, 3), padding='same', activation='relu')(x)
    x = layers.MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding='same')(x)

    x = TrajectoryGRU2D(filters=128, kernel_size=(3, 3), return_sequences=True)(x)
    x = layers.Conv3D(filters=128, kernel_size=(3, 3, 3), padding='same', activation='relu')(x)
    x = layers.MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding='same')(x)

    # Flatten before Fully Connected Layer
    x = layers.Flatten()(x)
    # x = layers.Dense(1, activation='sigmoid')(x)  

    model = models.Model(inputs=input_tensor, outputs=x)
    return model

def radial_structure_subnet(input_shape):
    """
    Creates the subnet for extracting TC radial structure features using a five-branch CNN design with 2D convolutions.

    Parameters:
    - input_shape: tuple, shape of the input data (e.g., (95, 95, 3))

    Returns:
    - model: tf.keras.Model, the radial structure subnet model
    """
    
    input_tensor = layers.Input(shape=input_shape)

    # Divide input data into four quadrants (NW, NE, SW, SE)
    # Assuming the input shape is (batch_size, height, width, channels)
    
    # Quadrant extraction - using slicing to separate quadrants
    nw_quadrant = input_tensor[:, :input_shape[0]//2, :input_shape[1]//2, :]
    ne_quadrant = input_tensor[:, :input_shape[0]//2, input_shape[1]//2:, :]
    sw_quadrant = input_tensor[:, input_shape[0]//2:, :input_shape[1]//2, :]
    se_quadrant = input_tensor[:, input_shape[0]//2:, input_shape[1]//2:, :]


    target_height = max(input_shape[0]//2, input_shape[0] - input_shape[0]//2)  # 48
    target_width = max(input_shape[1]//2, input_shape[1] - input_shape[1]//2)  # 48
    
    # Padding the quadrants to match the target size (48, 48)
    nw_quadrant = layers.ZeroPadding2D(padding=((0, target_height - nw_quadrant.shape[1]), 
                                                (0, target_width - nw_quadrant.shape[2])))(nw_quadrant)
    ne_quadrant = layers.ZeroPadding2D(padding=((0, target_height - ne_quadrant.shape[1]), 
                                                (0, target_width - ne_quadrant.shape[2])))(ne_quadrant)
    sw_quadrant = layers.ZeroPadding2D(padding=((0, target_height - sw_quadrant.shape[1]), 
                                                (0, target_width - sw_quadrant.shape[2])))(sw_quadrant)
    se_quadrant = layers.ZeroPadding2D(padding=((0, target_height - se_quadrant.shape[1]), 
                                                (0, target_width - se_quadrant.shape[2])))(se_quadrant)

    print(nw_quadrant.shape)
    print(ne_quadrant.shape)
    print(sw_quadrant.shape)
    print(se_quadrant.shape)
    # Main branch (processing the entire structure)
    main_branch = layers.Conv2D(filters=8, kernel_size=(3, 3), padding='same', activation='relu')(input_tensor)
    y=layers.MaxPool2D()(main_branch)

    y = layers.ZeroPadding2D(padding=((0, target_height - y.shape[1]), 
                                   (0, target_width - y.shape[2])))(y)
    # Side branches (processing the individual quadrants)
    nw_branch = layers.Conv2D(filters=8, kernel_size=(3, 3), padding='same', activation='relu')(nw_quadrant)
    ne_branch = layers.Conv2D(filters=8, kernel_size=(3, 3), padding='same', activation='relu')(ne_quadrant)
    sw_branch = layers.Conv2D(filters=8, kernel_size=(3, 3), padding='same', activation='relu')(sw_quadrant)
    se_branch = layers.Conv2D(filters=8, kernel_size=(3, 3), padding='same', activation='relu')(se_quadrant)
    
    # Apply padding to the side branches to match the dimensions of the main branch
    # nw_branch = layers.UpSampling2D(size=(2, 2), interpolation='nearest')(nw_branch)
    # ne_branch = layers.UpSampling2D(size=(2, 2), interpolation='nearest')(ne_branch)
    # sw_branch = layers.UpSampling2D(size=(2, 2), interpolation='nearest')(sw_branch)
    # se_branch = layers.UpSampling2D(size=(2, 2), interpolation='nearest')(se_branch)
    
    # Fusion operations (concatenate the outputs from the main branch and side branches)
    fusion = layers.concatenate([y, nw_branch, ne_branch, sw_branch, se_branch], axis=-1)
    
    # Additional convolution layer to combine the fused features
    x = layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation='relu')(fusion)
    x=layers.MaxPool2D(pool_size=(2, 2))(x)
    # Final dense layer for further processing
    nw_branch = layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation='relu')(nw_branch)
    
    ne_branch = layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation='relu')(ne_branch)
    sw_branch = layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation='relu')(sw_branch)
    se_branch = layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation='relu')(se_branch)
    nw_branch = layers.MaxPool2D(pool_size=(2, 2))(nw_branch)
    ne_branch = layers.MaxPool2D(pool_size=(2, 2))(ne_branch)
    sw_branch = layers.MaxPool2D(pool_size=(2, 2))(sw_branch)
    se_branch = layers.MaxPool2D(pool_size=(2, 2))(se_branch)

    fusion = layers.concatenate([x, nw_branch, ne_branch, sw_branch, se_branch], axis=-1)
    x = layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation='relu')(fusion)
    x=layers.MaxPool2D(pool_size=(2, 2))(x)
    
    nw_branch = layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation='relu')(nw_branch)
    
    ne_branch = layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation='relu')(ne_branch)
    sw_branch = layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation='relu')(sw_branch)
    se_branch = layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation='relu')(se_branch)
    nw_branch = layers.MaxPool2D(pool_size=(2, 2))(nw_branch)
    ne_branch = layers.MaxPool2D(pool_size=(2, 2))(ne_branch)
    sw_branch = layers.MaxPool2D(pool_size=(2, 2))(sw_branch)
    se_branch = layers.MaxPool2D(pool_size=(2, 2))(se_branch)

    fusion = layers.concatenate([x, nw_branch, ne_branch, sw_branch, se_branch], axis=-1)
    x = layers.Conv2D(filters=32, kernel_size=(3, 3),  activation='relu')(fusion)
    x=layers.Conv2D(filters=32, kernel_size=(3, 3), activation=None)(x)
    # Create and return the model
    x=layers.Flatten()(x)
    model = models.Model(inputs=input_tensor, outputs=x)
    return model

# Define input shape (batch_size, height, width, channels)
# input_shape = (95, 95, 8)  # Example input shape (95x95 spatial resolution, 3 channels)

# # Build the model
# model = radial_structure_subnet(input_shape)

# # Model summary
# model.summary()

def build_cnn_model(input_shape=(8, 8, 1)):
    # Define the input layer
    input_tensor = layers.Input(shape=input_shape)
    
    # Convolutional layer
    x = layers.Conv2D(64, (3, 3), padding='same')(input_tensor)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    
    # Flatten layer
    x = layers.Flatten()(x)
    
    # Create the model
    model = models.Model(inputs=input_tensor, outputs=x)
    
    return model

from tensorflow.keras import layers, models, Input # type: ignore

def build_combined_model():
    # Define input shapes
    input_shape_3d = (8, 95, 95, 2)
    input_shape_radial = (95, 95, 8)
    input_shape_cnn = (8, 8, 1)
    
    input_shape_latitude = (8,)
    input_shape_longitude = (8,)
    input_shape_other = (9,)

    # Build individual models
    model_3d = build_tgru_model(input_shape=input_shape_3d)
    model_radial = radial_structure_subnet(input_shape=input_shape_radial)
    model_cnn = build_cnn_model(input_shape=input_shape_cnn)

    # Define new inputs
    input_latitude = Input(shape=input_shape_latitude ,name="latitude_input")
    input_longitude = Input(shape=input_shape_longitude, name="longitude_input")
    input_other = Input(shape=input_shape_other, name="other_input")

    # Flatten the additional inputs
    flat_latitude = layers.Dense(32,activation='relu')(input_latitude)
    flat_longitude = layers.Dense(32,activation='relu')(input_longitude)
    flat_other = layers.Dense(64,activation='relu')(input_other)

    # Combine all outputs
    combined = layers.concatenate([
        model_3d.output, 
        model_radial.output, 
        model_cnn.output,
        flat_latitude, 
        flat_longitude, 
        flat_other
    ])

    # Add dense layers for final processing
    x = layers.Dense(128, activation='relu')(combined)  
    x = layers.Dense(1, activation=None)(x)

    # Create the final model
    final_model = models.Model(
        inputs=[model_3d.input, model_radial.input, model_cnn.input,
                input_latitude, input_longitude, input_other ],
        outputs=x
    )

    return final_model

import h5py
# with h5py.File(r"Trj_GRU.h5", 'r') as f:
#     print(f.attrs.get('keras_version'))
#     print(f.attrs.get('backend'))

#     print("Model layers:", list(f['model_weights'].keys()))

model = build_combined_model()  # Your original model building function
# Rebuild the model architecture
# Step 1: Build the full combined model (with 6 inputs)
# model = build_combined_model()

# Step 2: Call the model once with dummy data to build the weights
# import tensorflow as tf

dummy_input = [
    tf.random.normal((1, 8, 95, 95, 2)),   # reduced_images_test
    tf.random.normal((1, 95, 95, 8)),      # hov_m_test
    tf.random.normal((1, 8, 8, 1)),        # test_vmax_3d
    tf.random.normal((1, 8)),              # lat_test
    tf.random.normal((1, 8)),              # lon_test
    tf.random.normal((1, 9)),              # other_scalar_inputs
]
_ = model(dummy_input)  # Build model by doing one forward pass

# Step 3: Load weights
model.load_weights("Trj_GRU.weights.h5")  # Make sure this matches the architecture



def predict_trajgru(reduced_images_test,hov_m_test,test_vmax_3d,lat_test,lon_test,int_diff_test):
    y=model.predict([reduced_images_test,hov_m_test,test_vmax_3d,lat_test,lon_test,int_diff_test ])
    return y