Troubles with training rotation models

[jwutsetro]

Dear AshStuff,

Thank you for this very complete GitHub repo. At our research group, we are very interested to use dismm models to predict some anatomical structures on both MRI and CT images for multiple projects. For that, we would want to train and infer models that are trained with pca loads.

Unfortunately we are stuck after training the scale models. We manage to update our train and val jsons after the initial translation only training. But hen we run into some errors related to config files when trying to train with rotation and PCA. They are related with the global_init_mtx argument in the transform pipeline for rotation and PCA. How should we initialise these values based on our trained translate and scale models. Some hardcoded examples are given for the liver and the larynx, but we don't know how this translates for our custom datasets.

We tried giving some median values from the jsons to construct it
`

class ConfigFromJSON:
    def __init__(self, json_path):
    self.json_path = json_path
    self.trans_center = None
    self.scale = None
    self.Apx2sdf = None
    self.global_init_mtx = None
    self.compute_parameters()

def compute_parameters(self):
    with open(self.json_path, 'r') as file:
        data = json.load(file)
    
    # Extract translation vectors, rotation matrices, and scaling factors
    translations = np.array([entry['t'] for entry in data])
    rotations = [np.array(entry['R']) for entry in data]
    scales = np.array([entry['s'] for entry in data])

    # Compute median translation and mean scale
    self.trans_center = np.median(translations, axis=0)
    self.scale = np.mean(scales)

    # Compute mean rotation matrix using SVD (simplified approach)
    mean_rotation = np.mean(np.stack(rotations), axis=0)
    u, _, v = np.linalg.svd(mean_rotation)
    self.Apx2sdf = np.dot(u, v)

    # Scale and translate in the affine matrix
    self.Apx2sdf = np.diag([1/self.scale, 1/self.scale, 1/self.scale, 1])
    self.Apx2sdf[:3, 3] = -self.trans_center

    # Optional: Compute a global initial matrix
    self.global_init_mtx = np.eye(4)
    self.global_init_mtx[:3, 3] = np.mean(translations, axis=0)

def __repr__(self):
    return (f"CustomConfig(trans_center={self.trans_center}, scale={self.scale}, "
            f"Apx2sdf={self.Apx2sdf}, global_init_mtx={self.global_init_mtx})")

Even though we are not sure if this works, we immediately run into an other error related to a predict_scale parameter that is missing in the json dict when training with rotation. Do you perhaps have a working example of the rotation and pca training?

Thanks in advance,
J

`