Update prediction.py

prediction.py

@@ -149,12 +149,12 @@ def both_joint_loss_function(y_true, y_pred):
     return detection_loss + pp_loss
 
 
+import h5py
+import json
+from keras.models import load_model
+
 def load_old_model(model_file):
-    import h5py
-    from keras.models import load_model
-    from keras.optimizers import Adam  # or the relevant optimizer you expect
     print("Loading pre-trained model")
-
     custom_objects = {
         'mse': mse,
         'weighted_mean_se': weighted_mean_se,
@@ -174,18 +174,22 @@ def load_old_model(model_file):
         pass
 
     try:
-        # Intercept and fix optimizer config before loading
-        with h5py.File(model_file, 'r') as f:
+        # Patch 'lr' to 'learning_rate' if necessary
+        with h5py.File(model_file, 'r+') as f:
             if 'training_config' in f.attrs:
-                import json
-                training_config = json.loads(f.attrs['training_config'].decode('utf-8'))
-                optimizer_config = training_config['optimizer_config']
+                training_config_json = f.attrs['training_config']
+                training_config = json.loads(training_config_json)  # <- FIXED
+
+                optimizer_config = training_config.get('optimizer_config', {})
                 config = optimizer_config.get('config', {})
 
-                # Replace deprecated 'lr' with 'learning_rate'
                 if 'lr' in config:
                     config['learning_rate'] = config.pop('lr')
                     optimizer_config['config'] = config
+                    training_config['optimizer_config'] = optimizer_config
+
+                    # Update the HDF5 file with modified config
+                    f.attrs.modify('training_config', json.dumps(training_config))
 
         return load_model(model_file, custom_objects=custom_objects)