Skip to content

feat: Retrain Attention ResUNet segmentation model and resolve dimension mismatch#93

Open
Nicks-19 wants to merge 2 commits into
Brijeshthummar02:masterfrom
Nicks-19:feat/retrain-segmentation-model
Open

feat: Retrain Attention ResUNet segmentation model and resolve dimension mismatch#93
Nicks-19 wants to merge 2 commits into
Brijeshthummar02:masterfrom
Nicks-19:feat/retrain-segmentation-model

Conversation

@Nicks-19

@Nicks-19 Nicks-19 commented Jun 6, 2026

Copy link
Copy Markdown

Overview

This Pull Request updates the 162-layer Attention ResUNet brain tumor segmentation model. It uses a dataset and it fixes some problems, with the Flask backend. This fixes issue #74.

Key Modifications

  1. Model Architecture & Compatibility Updates (app.py)

The code now looks for ResUNet-weights.keras first when it loads the segmentation model. If it does not find this it will use weights_seg.hdf5 instead.
This fixes a problem where the wrong number of layers were loaded. This caused an error that said ValueError: Layer count mismatch....

  1. Loss & Metrics Precision Fixes (utilities.py)

The code now makes sure that y_true and y_pred are the type when it calculates the loss and metrics. This prevents a TypeError from happening when the code is run.
This is because Keras 3 and TensorFlow 2.x need float32 but the code was giving them int64.

  1. Standalone Training Pipeline (train_segmentation.py. ResUNet-model.json`)

There is a Python script that can train the ResUNet model on its own. It builds the model trains it and saves the weights.
The script can train both the 90-layer simple ResUNet and the 162-layer Attention ResUNet.
The ResUNet-model.json file has been updated to make sure everything is consistent.

  1. Model Ignored in Git (.gitignore)

The .gitignore file now tells Git to ignore *.keras and *.h5 files. These are the weight files that we do not want to commit to GitHub.


Verification and Metrics

1. Flask App Initialization

Verified that the backend starts up and loads all ensemble and segmentation models successfully without any dimensions mismatch:

Loading primary classification model...
✓ Primary classification model loaded successfully
Loading secondary classification model...
✓ Secondary classification model loaded with JSON architecture
Loading primary segmentation model (ResUNet-MRI)...
✓ Primary segmentation model loaded successfully
Loading secondary segmentation model...
✓ Secondary segmentation model loaded successfully
    Model Summary:
   - Classification models loaded: 2
   - Segmentation models loaded: 2
   - TTA enabled: True
   - Ensemble enabled: True
All models loaded successfully: True

@Brijeshthummar02

Copy link
Copy Markdown
Owner

@Nicks-19 add issue no in pr description

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants