Skip to content

Commit 2c3d44c

Browse files
authored
🐛 Fix: Memory leak in image processing endpoint (#1513)
1 parent 5a0d662 commit 2c3d44c

File tree

1 file changed

+6
-4
lines changed

1 file changed

+6
-4
lines changed

Diff for: examples/sam2_amg_server/server.py

+6-4
Original file line numberDiff line numberDiff line change
@@ -555,19 +555,21 @@ async def upload_image(image: UploadFile = File(...)):
555555
response_future = asyncio.Future()
556556
await request_queue.put((image_tensor, response_future))
557557
masks = await response_future
558-
559-
# Save an example
560-
plt.figure(figsize=(image_tensor.shape[1]/100., image_tensor.shape[0]/100.), dpi=100)
558+
559+
# Create figure and ensure it's closed after generating response
560+
fig = plt.figure(figsize=(image_tensor.shape[1]/100., image_tensor.shape[0]/100.), dpi=100)
561561
plt.imshow(image_tensor)
562562
show_anns(masks, rle_to_mask)
563563
plt.axis('off')
564564
plt.tight_layout()
565+
565566
buf = BytesIO()
566567
plt.savefig(buf, format='png')
567568
buf.seek(0)
569+
plt.close(fig) # Close figure after we're done with it
570+
568571
return StreamingResponse(buf, media_type="image/png")
569572

570-
571573
# uvicorn.run(app, host=host, port=port, log_level="info")
572574
uvicorn.run(app, host=host, port=port)
573575

0 commit comments

Comments
 (0)