Update main.py

main.py

@@ -32,21 +32,21 @@ async def segment(files: list[UploadFile]):
         #     return FileResponse(tempFile.name, media_type="multipart/form-data")
         
         with io.BytesIO() as buffer:
-        # Save the NIfTI image to the 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())
+                    gz_file.write(buffer.getvalue())  # Write the buffer content directly
+        
+                # Reset the buffer cursor to the beginning
+                gz_buffer.seek(0)
+        
+                # Return the compressed content as a streaming response
+                return StreamingResponse(content=gz_buffer, media_type="application/octet-stream")
         
-                # 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)}")
         return {"msg": "An error occurred during segmentation."}