mathpluscode commited on
Commit
2e9f7e2
Β·
1 Parent(s): 7ca35bb

Add results cache

Browse files
Files changed (2) hide show
  1. app.py +45 -19
  2. requirements.txt +1 -1
app.py CHANGED
@@ -176,6 +176,14 @@ def mae_inference(
176
 
177
 
178
  def mae(image_id, mask_ratio, progress=gr.Progress()):
 
 
 
 
 
 
 
 
179
  t = 4 # which time frame to use
180
  progress(0, desc="Downloading model...")
181
  model = CineMA.from_pretrained()
@@ -221,7 +229,6 @@ def mae(image_id, mask_ratio, progress=gr.Progress()):
221
  masks_dict["sax"] = np.transpose(masks_dict["sax"], (1, 0, 2))
222
 
223
  # Plot MAE reconstruction and save to file
224
- mae_path = cache_dir / f"mae_image{image_id}_mask{mask_ratio:.2f}.png"
225
  plot_mae_reconstruction(batch, reconstructed_dict, masks_dict, mae_path)
226
 
227
  return str(mae_path)
@@ -298,6 +305,21 @@ def segmentation_sax_inference(
298
 
299
 
300
  def segmentation_sax(trained_dataset, seed, image_id, t_step, progress=gr.Progress()):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
301
  # Fixed parameters
302
  view = "sax"
303
  split = "train" if image_id <= 100 else "test"
@@ -343,18 +365,8 @@ def segmentation_sax(trained_dataset, seed, image_id, t_step, progress=gr.Progre
343
 
344
  progress(1, desc="Inference finished. Plotting ...")
345
 
346
- # Create file paths for saving plots
347
- seg_path = (
348
- cache_dir
349
- / f"sax_segmentation_{trained_dataset}_image{image_id}_seed{seed}_tstep{t_step}.gif"
350
- )
351
- vol_path = (
352
- cache_dir
353
- / f"sax_volume_{trained_dataset}_image{image_id}_seed{seed}_tstep{t_step}.png"
354
- )
355
-
356
  # Plot segmentations and volume changes with file paths
357
- plot_segmentations_sax(images, labels, seg_path)
358
  plot_volume_changes_sax(labels, t_step, vol_path)
359
 
360
  return (str(seg_path), str(vol_path))
@@ -474,6 +486,15 @@ def segmentation_lax_inference(
474
 
475
 
476
  def segmentation_lax(seed, image_id, progress=gr.Progress()):
 
 
 
 
 
 
 
 
 
477
  # Fixed parameters
478
  trained_dataset = "mnms2"
479
  view = "lax_4c"
@@ -502,11 +523,9 @@ def segmentation_lax(seed, image_id, progress=gr.Progress()):
502
  progress(1, desc="Inference finished. Plotting ...")
503
 
504
  # Plot segmentations and save as GIF
505
- seg_path = cache_dir / f"lax_segmentation_image{image_id}_seed{seed}.gif"
506
  plot_segmentations_lax(images, labels, seg_path)
507
 
508
  # Plot volume changes and save as figure
509
- vol_path = cache_dir / f"lax_volume_image{image_id}_seed{seed}.png"
510
  plot_volume_changes_lax(labels, vol_path)
511
 
512
  return (str(seg_path), str(vol_path))
@@ -647,6 +666,17 @@ def landmark(image_id, view, method, seed, progress=gr.Progress()):
647
  view = "lax_2c" if view == "LAX 2C" else "lax_4c"
648
  method = method.lower()
649
 
 
 
 
 
 
 
 
 
 
 
 
650
  # Download and load model
651
  progress(0, desc="Downloading model...")
652
  if method == "heatmap":
@@ -686,13 +716,9 @@ def landmark(image_id, view, method, seed, progress=gr.Progress()):
686
  progress(1, desc="Inference finished. Plotting ...")
687
 
688
  # Plot landmarks in GIF
689
- landmark_path = (
690
- cache_dir / f"landmark_{view}_image{image_id}_{method}_seed{seed}.gif"
691
- )
692
  plot_landmarks(images, coords, landmark_path)
693
 
694
  # Plot LV change in PNG
695
- lv_path = cache_dir / f"lv_{view}_image{image_id}_{method}_seed{seed}.png"
696
  plot_lv(coords, lv_path)
697
 
698
  return (str(landmark_path), str(lv_path))
@@ -791,7 +817,7 @@ with gr.Blocks(
791
  # CineMA: A Foundation Model for Cine Cardiac MRI πŸŽ₯πŸ«€
792
 
793
  πŸš€ The following demos showcase the capabilities of CineMA in multiple tasks.<br>
794
- ⏱️ The examples may take 10-60 seconds to download data and model, perform inference, and render plots.<br>
795
  πŸ”— For more details, check out our [GitHub](https://github.com/mathpluscode/CineMA).
796
  """
797
  )
 
176
 
177
 
178
  def mae(image_id, mask_ratio, progress=gr.Progress()):
179
+ # Create file path for saving MAE reconstruction plot
180
+ mae_path = cache_dir / f"mae_image{image_id}_mask{mask_ratio * 100:.0f}.png"
181
+
182
+ # Check if result already exists
183
+ if mae_path.exists():
184
+ progress(1, desc="Loading cached result...")
185
+ return str(mae_path)
186
+
187
  t = 4 # which time frame to use
188
  progress(0, desc="Downloading model...")
189
  model = CineMA.from_pretrained()
 
229
  masks_dict["sax"] = np.transpose(masks_dict["sax"], (1, 0, 2))
230
 
231
  # Plot MAE reconstruction and save to file
 
232
  plot_mae_reconstruction(batch, reconstructed_dict, masks_dict, mae_path)
233
 
234
  return str(mae_path)
 
305
 
306
 
307
  def segmentation_sax(trained_dataset, seed, image_id, t_step, progress=gr.Progress()):
308
+ # Create file paths for saving plots
309
+ seg_path = (
310
+ cache_dir
311
+ / f"sax_segmentation_{trained_dataset}_image{image_id}_seed{seed}_tstep{t_step}.gif"
312
+ )
313
+ vol_path = (
314
+ cache_dir
315
+ / f"sax_volume_{trained_dataset}_image{image_id}_seed{seed}_tstep{t_step}.png"
316
+ )
317
+
318
+ # Check if results already exist
319
+ if seg_path.exists() and vol_path.exists():
320
+ progress(1, desc="Loading cached results...")
321
+ return (str(seg_path), str(vol_path))
322
+
323
  # Fixed parameters
324
  view = "sax"
325
  split = "train" if image_id <= 100 else "test"
 
365
 
366
  progress(1, desc="Inference finished. Plotting ...")
367
 
 
 
 
 
 
 
 
 
 
 
368
  # Plot segmentations and volume changes with file paths
369
+ plot_segmentations_sax(images, labels, t_step, seg_path)
370
  plot_volume_changes_sax(labels, t_step, vol_path)
371
 
372
  return (str(seg_path), str(vol_path))
 
486
 
487
 
488
  def segmentation_lax(seed, image_id, progress=gr.Progress()):
489
+ # Create file paths for saving plots
490
+ seg_path = cache_dir / f"lax_segmentation_image{image_id}_seed{seed}.gif"
491
+ vol_path = cache_dir / f"lax_volume_image{image_id}_seed{seed}.png"
492
+
493
+ # Check if results already exist
494
+ if seg_path.exists() and vol_path.exists():
495
+ progress(1, desc="Loading cached results...")
496
+ return (str(seg_path), str(vol_path))
497
+
498
  # Fixed parameters
499
  trained_dataset = "mnms2"
500
  view = "lax_4c"
 
523
  progress(1, desc="Inference finished. Plotting ...")
524
 
525
  # Plot segmentations and save as GIF
 
526
  plot_segmentations_lax(images, labels, seg_path)
527
 
528
  # Plot volume changes and save as figure
 
529
  plot_volume_changes_lax(labels, vol_path)
530
 
531
  return (str(seg_path), str(vol_path))
 
666
  view = "lax_2c" if view == "LAX 2C" else "lax_4c"
667
  method = method.lower()
668
 
669
+ # Create file paths for saving plots
670
+ landmark_path = (
671
+ cache_dir / f"landmark_{view}_image{image_id}_{method}_seed{seed}.gif"
672
+ )
673
+ lv_path = cache_dir / f"lv_{view}_image{image_id}_{method}_seed{seed}.png"
674
+
675
+ # Check if results already exist
676
+ if landmark_path.exists() and lv_path.exists():
677
+ progress(1, desc="Loading cached results...")
678
+ return (str(landmark_path), str(lv_path))
679
+
680
  # Download and load model
681
  progress(0, desc="Downloading model...")
682
  if method == "heatmap":
 
716
  progress(1, desc="Inference finished. Plotting ...")
717
 
718
  # Plot landmarks in GIF
 
 
 
719
  plot_landmarks(images, coords, landmark_path)
720
 
721
  # Plot LV change in PNG
 
722
  plot_lv(coords, lv_path)
723
 
724
  return (str(landmark_path), str(lv_path))
 
817
  # CineMA: A Foundation Model for Cine Cardiac MRI πŸŽ₯πŸ«€
818
 
819
  πŸš€ The following demos showcase the capabilities of CineMA in multiple tasks.<br>
820
+ ⏱️ The examples may take 10-60 seconds, if not cached, to download data and model, perform inference, and render plots.<br>
821
  πŸ”— For more details, check out our [GitHub](https://github.com/mathpluscode/CineMA).
822
  """
823
  )
requirements.txt CHANGED
@@ -17,6 +17,6 @@ scikit-learn==1.6.1
17
  scipy==1.15.2
18
  spaces==0.36.0
19
  timm==1.0.15
20
- git+https://github.com/mathpluscode/CineMA@7e86ffc7ddf06ad7283915ee143ed808c0f59576#egg=cinema
21
  --extra-index-url https://download.pytorch.org/whl/cu113
22
  torch==2.5.1
 
17
  scipy==1.15.2
18
  spaces==0.36.0
19
  timm==1.0.15
20
+ git+https://github.com/mathpluscode/CineMA@0d1afe864d4b4c348a993b8b2b790adf8581bc03#egg=cinema
21
  --extra-index-url https://download.pytorch.org/whl/cu113
22
  torch==2.5.1