Update main.py

main.py

@@ -23,9 +23,10 @@ async def segment(files: list[UploadFile]):
         print("input processed")
         output = segment_all(wt_model, tc_model, et_model, flair_tensor, t1_tensor, t1ce_tensor, t2_tensor)
         print("Segmented Successfuly!", output.shape)
-        
+        output_img = nib.Nifti1Image(output, affine=np.eye(4))
+
         with tempfile.NamedTemporaryFile(suffix=".nii") as tempFile:
-            nib.save(output, tempFile.name)
+            nib.save(output_img, tempFile.name)
             return FileResponse(tempFile.name, media_type="multipart/form-data")
     
     except Exception as e: