ayyuce commited on
Commit
2623207
·
verified ·
1 Parent(s): fad7d2c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -5
app.py CHANGED
@@ -13,9 +13,13 @@ def run_segmentation(uploaded_file, modality):
13
  input_filename = f"input_{job_id}.nii.gz"
14
  output_folder = f"segmentations_{job_id}"
15
 
16
-
17
- with open(input_filename, "wb") as f:
18
- f.write(uploaded_file.read())
 
 
 
 
19
 
20
  command = ["TotalSegmentator", "-i", input_filename, "-o", output_folder]
21
  if modality == "MR":
@@ -34,7 +38,6 @@ def run_segmentation(uploaded_file, modality):
34
  arcname = os.path.relpath(file_path, output_folder)
35
  zipf.write(file_path, arcname)
36
 
37
-
38
  seg_files = [os.path.join(output_folder, f) for f in os.listdir(output_folder) if f.endswith('.nii.gz')]
39
  if seg_files:
40
  seg_file = seg_files[0]
@@ -55,7 +58,6 @@ def run_segmentation(uploaded_file, modality):
55
  else:
56
  image_filename = None
57
 
58
-
59
  os.remove(input_filename)
60
  shutil.rmtree(output_folder)
61
 
 
13
  input_filename = f"input_{job_id}.nii.gz"
14
  output_folder = f"segmentations_{job_id}"
15
 
16
+ if isinstance(uploaded_file, str):
17
+ shutil.copy(uploaded_file, input_filename)
18
+ elif hasattr(uploaded_file, "read"):
19
+ with open(input_filename, "wb") as f:
20
+ f.write(uploaded_file.read())
21
+ else:
22
+ return "Invalid file input", None
23
 
24
  command = ["TotalSegmentator", "-i", input_filename, "-o", output_folder]
25
  if modality == "MR":
 
38
  arcname = os.path.relpath(file_path, output_folder)
39
  zipf.write(file_path, arcname)
40
 
 
41
  seg_files = [os.path.join(output_folder, f) for f in os.listdir(output_folder) if f.endswith('.nii.gz')]
42
  if seg_files:
43
  seg_file = seg_files[0]
 
58
  else:
59
  image_filename = None
60
 
 
61
  os.remove(input_filename)
62
  shutil.rmtree(output_folder)
63