darragh commited on
Commit
68facde
β€’
1 Parent(s): 0301a69

Test video output

Browse files
Files changed (4) hide show
  1. README.md +9 -4
  2. app.py +113 -0
  3. packages.txt +1 -0
  4. requirements.txt +12 -0
README.md CHANGED
@@ -1,8 +1,8 @@
1
  ---
2
  title: Swinunetr Dicom Video
3
- emoji: πŸƒ
4
- colorFrom: purple
5
- colorTo: red
6
  sdk: gradio
7
  sdk_version: 3.0.24
8
  app_file: app.py
@@ -10,4 +10,9 @@ pinned: false
10
  license: apache-2.0
11
  ---
12
 
13
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
1
  ---
2
  title: Swinunetr Dicom Video
3
+ emoji: πŸ“–πŸŽ¬
4
+ colorFrom: blue
5
+ colorTo: purple
6
  sdk: gradio
7
  sdk_version: 3.0.24
8
  app_file: app.py
 
10
  license: apache-2.0
11
  ---
12
 
13
+ This repository contains the code for UNETR: Transformers for 3D Medical Image Segmentation. UNETR is the first 3D segmentation network that uses a pure vision transformer as its encoder without relying on CNNs for feature extraction. The code presents a volumetric (3D) multi-organ segmentation application using the BTCV challenge dataset.
14
+
15
+ Check out the Beyond the Cranial Vault source Swin-UNET models [here](https://huggingface.co/darragh/swinunetr-btcv-small). Also in the link, you can see links to the original BTCV winning solution.
16
+
17
+ This is a small demo on a subset of the test data for the [BTCV competition](https://zenodo.org/record/1169361#.YtGvn-xKhb8).
18
+
app.py ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sys
2
+ import os
3
+ import glob
4
+ import shutil
5
+ import torch
6
+ import argparse
7
+ import mediapy
8
+ import cv2
9
+ import numpy as np
10
+ import gradio as gr
11
+ from skimage import color, img_as_ubyte
12
+ from monai import transforms, data
13
+
14
+ os.system("git clone https://github.com/darraghdog/Project-MONAI-research-contributions pmrc")
15
+ sys.path.append("pmrc/SwinUNETR/BTCV")
16
+ from swinunetr import SwinUnetrModelForInference, SwinUnetrConfig
17
+
18
+
19
+ ffmpeg_path = shutil.which('ffmpeg')
20
+ mediapy.set_ffmpeg(ffmpeg_path)
21
+
22
+ model = SwinUnetrModelForInference.from_pretrained('darragh/swinunetr-btcv-tiny')
23
+ model.eval()
24
+
25
+ input_files = glob.glob('pmrc/SwinUNETR/BTCV/dataset/imagesSampleTs/*.nii.gz')
26
+ input_files = dict((f.split('/')[-1], f) for f in input_files)
27
+
28
+ # Load and process dicom with monai transforms
29
+ test_transform = transforms.Compose(
30
+ [
31
+ transforms.LoadImaged(keys=["image"]),
32
+ transforms.AddChanneld(keys=["image"]),
33
+ transforms.Spacingd(keys="image",
34
+ pixdim=(1.5, 1.5, 2.0),
35
+ mode="bilinear"),
36
+ transforms.ScaleIntensityRanged(keys=["image"],
37
+ a_min=-175.0,
38
+ a_max=250.0,
39
+ b_min=0.0,
40
+ b_max=1.0,
41
+ clip=True),
42
+ # transforms.Resized(keys=["image"], spatial_size = (256,256,-1)),
43
+ transforms.ToTensord(keys=["image"]),
44
+ ])
45
+
46
+
47
+
48
+ # Create Data Loader
49
+ def create_dl(test_files):
50
+ ds = test_transform(test_files)
51
+ loader = data.DataLoader(ds,
52
+ batch_size=1,
53
+ shuffle=False)
54
+ return loader
55
+
56
+ # Inference and video generation
57
+ def generate_dicom_video(selected_file):
58
+
59
+ test_file = input_files[selected_file]
60
+ test_files = [{'image': test_file}]
61
+ dl = create_dl(test_files)
62
+ batch = next(iter(dl))
63
+
64
+ tst_inputs = batch["image"]
65
+ tst_inputs = tst_inputs[:,:,:,:,-32:]
66
+
67
+ with torch.no_grad():
68
+ outputs = model(tst_inputs,
69
+ (96,96,96),
70
+ 8,
71
+ overlap=0.5,
72
+ mode="gaussian")
73
+ tst_outputs = torch.softmax(outputs.logits, 1)
74
+ tst_outputs = torch.argmax(tst_outputs, axis=1)
75
+
76
+ # Write frames to video
77
+ for inp, outp in zip(tst_inputs, tst_outputs):
78
+
79
+ frames = []
80
+ for idx in range(inp.shape[-1]):
81
+ # Segmentation
82
+ seg = outp[:,:,idx].numpy().astype(np.uint8)
83
+ # Input dicom frame
84
+ img = (inp[0,:,:,idx]*255).numpy().astype(np.uint8)
85
+ img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)
86
+ frame = color.label2rgb(seg,img, bg_label = 0)
87
+ frame = img_as_ubyte(frame)
88
+ frame = np.concatenate((img, frame), 1)
89
+ frames.append(frame)
90
+ mediapy.write_video("dicom.mp4", frames, fps=4)
91
+
92
+ return 'dicom.mp4'
93
+
94
+
95
+ '''
96
+ test_file = glob.glob('pmrc/SwinUNETR/BTCV/dataset/imagesSampleTs/*.nii.gz')[0]
97
+ generate_dicom_video(test_file)
98
+ '''
99
+
100
+ demo = gr.Blocks()
101
+
102
+ with demo:
103
+ selected_dicom_key = gr.inputs.Dropdown(
104
+ choices=sorted(input_files),
105
+ type="value",
106
+ label="Select a dicom file")
107
+ button_gen_video = gr.Button("Generate Video")
108
+ output_interpolation = gr.Video(label="Generated Video")
109
+ button_gen_video.click(fn=generate_dicom_video, inputs=selected_dicom_key, outputs=output_interpolation)
110
+
111
+ demo.launch(debug=True, enable_queue=True)
112
+
113
+
packages.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ ffmpeg
requirements.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ transformers==4.20.1
2
+ torch==1.10.0
3
+
4
+ git+https://github.com/Project-MONAI/MONAI#egg.gitmonai@0.8.1+271.g07de215c
5
+ nibabel==3.1.1
6
+ tqdm==4.59.0
7
+ einops==0.4.1
8
+ tensorboardX==2.1
9
+ scipy==1.5.0
10
+ mediapy==1.0.3
11
+ scikit-image==0.17.2
12
+ opencv-python==4.6.0.66