ItzRoBeerT commited on
Commit
2ccc1a1
·
1 Parent(s): 17de1f9

Remove image background

Browse files
Files changed (3) hide show
  1. .gitignore +1 -1
  2. requirements.txt +4 -0
  3. src/utils/image_helpers.py +58 -0
.gitignore CHANGED
@@ -1 +1 @@
1
- __pycache__/*
 
1
+ __pycache__/
requirements.txt CHANGED
@@ -1 +1,5 @@
1
  fastmcp
 
 
 
 
 
1
  fastmcp
2
+ requests
3
+ Pillow
4
+ rembg
5
+ onnxruntime
src/utils/image_helpers.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import requests
2
+ from typing import Optional, Dict, Any
3
+ import os
4
+ import rembg
5
+
6
+ async def remove_background_from_url(
7
+ image_url: str,
8
+ output_path: str,
9
+ model_name: str = "u2net"
10
+ ) -> Dict[str, Any]:
11
+ """
12
+ Remove background from an image downloaded from a URL.
13
+
14
+ Args:
15
+ image_url: URL of the image to process
16
+ output_path: Path where to save the processed image
17
+ model_name: Background removal model to use
18
+
19
+ Returns:
20
+ Dictionary with result information
21
+ """
22
+
23
+ try:
24
+ # Download image from URL
25
+ response = requests.get(image_url, timeout=30)
26
+ response.raise_for_status()
27
+
28
+ # Remove background
29
+ session = rembg.new_session(model_name=model_name)
30
+ output_data = rembg.remove(response.content, session=session)
31
+
32
+ # Save the result
33
+ with open(output_path, 'wb') as output_file:
34
+ output_file.write(output_data)
35
+
36
+ output_size = os.path.getsize(output_path)
37
+
38
+ return {
39
+ "success": True,
40
+ "message": f"Background removed from URL image using {model_name} model",
41
+ "source_url": image_url,
42
+ "output_path": output_path,
43
+ "output_size_bytes": output_size,
44
+ "model_used": model_name
45
+ }
46
+
47
+ except requests.RequestException as e:
48
+ return {
49
+ "success": False,
50
+ "error": f"Failed to download image: {str(e)}",
51
+ "output_path": None
52
+ }
53
+ except Exception as e:
54
+ return {
55
+ "success": False,
56
+ "error": f"Failed to process image: {str(e)}",
57
+ "output_path": None
58
+ }