Update main.py

main.py

@@ -1,6 +1,7 @@
 from model_loader import load_vChain_model
 from inference import preprocess_input, segment_all
-
+import io
+import gzip
 import tempfile
 import numpy as np
 import nibabel as nib
@@ -26,9 +27,25 @@ async def segment(files: list[UploadFile]):
         print("Segmented Successfuly!", output.shape)
         output_img = nib.Nifti1Image(output, affine=np.eye(4))
 
-        with tempfile.NamedTemporaryFile(suffix=".nii") as tempFile:
-            nib.save(output_img, tempFile.name)
-            return FileResponse(tempFile.name, media_type="multipart/form-data")
+        # with tempfile.NamedTemporaryFile(suffix=".nii") as tempFile:
+        #     nib.save(output_img, tempFile.name)
+        #     return FileResponse(tempFile.name, media_type="multipart/form-data")
+        
+        with io.BytesIO() as buffer:
+        # Save the NIfTI image to the buffer
+        nib.save(nifti_img, buffer)
+        buffer.seek(0)  # Reset the buffer cursor to the beginning
+    
+        # Compress the buffer content using gzip
+        with io.BytesIO() as gz_buffer:
+            with gzip.GzipFile(fileobj=gz_buffer, mode="wb") as gz_file:
+                gz_file.write(buffer.read())
+    
+            # Get the compressed content
+            compressed_content = gz_buffer.getvalue()
+    
+        # Return the compressed content as a streaming response
+        return StreamingResponse(content=compressed_content, media_type="application/octet-stream")
     
     except Exception as e:
         print(f"Error occurred: {str(e)}")