JuanjoSG5 commited on
Commit
b2b7174
·
2 Parent(s): 2d2877d 195f21b

Merge branch 'main' of https://huggingface.co/spaces/AgentsGuards/agents-guard-mcp

Browse files
mcp_server.py CHANGED
@@ -1,34 +1,28 @@
1
  from mcp.server.fastmcp import FastMCP
2
  from src.utils.change_format import change_format
3
- from src.utils.image_helpers import remove_background_from_url
4
  from src.utils.visualize_image import visualize_base64_image
5
  from src.utils.generate_image import generate_image
6
  from src.utils.apply_filter import apply_filter
7
  from src.utils.add_text import add_text_to_image
8
  from src.utils.resize_image import resize_image
 
 
 
9
 
10
  mcp = FastMCP("Youtube Service")
11
 
12
- @mcp.tool()
13
- def say_hello(name: str) -> str:
14
- """
15
- Returns a greeting message for the given name.
16
-
17
- Args:
18
- name (str): The name to greet.
19
-
20
- Returns:
21
- str: A greeting message.
22
- """
23
- return f"Hello, {name}!"
24
-
25
  mcp.add_tool(remove_background_from_url)
 
26
  mcp.add_tool(change_format)
27
  mcp.add_tool(visualize_base64_image)
28
  mcp.add_tool(generate_image)
29
  mcp.add_tool(apply_filter)
30
  mcp.add_tool(add_text_to_image)
31
  mcp.add_tool(resize_image)
 
 
 
32
 
33
  if __name__ == "__main__":
34
  mcp.run()
 
1
  from mcp.server.fastmcp import FastMCP
2
  from src.utils.change_format import change_format
3
+ from src.utils.remove_background import remove_background_from_url
4
  from src.utils.visualize_image import visualize_base64_image
5
  from src.utils.generate_image import generate_image
6
  from src.utils.apply_filter import apply_filter
7
  from src.utils.add_text import add_text_to_image
8
  from src.utils.resize_image import resize_image
9
+ from src.utils.watermark import add_watermark, remove_watermark
10
+ from src.utils.describe import describe_image
11
+ from src.utils.compress import compress_image
12
 
13
  mcp = FastMCP("Youtube Service")
14
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
  mcp.add_tool(remove_background_from_url)
16
+ mcp.add_tool(describe_image)
17
  mcp.add_tool(change_format)
18
  mcp.add_tool(visualize_base64_image)
19
  mcp.add_tool(generate_image)
20
  mcp.add_tool(apply_filter)
21
  mcp.add_tool(add_text_to_image)
22
  mcp.add_tool(resize_image)
23
+ mcp.add_tool(add_watermark)
24
+ mcp.add_tool(remove_watermark)
25
+ mcp.add_tool(compress_image)
26
 
27
  if __name__ == "__main__":
28
  mcp.run()
requirements.txt CHANGED
@@ -3,4 +3,5 @@ requests
3
  Pillow
4
  rembg
5
  onnxruntime
6
- openai
 
 
3
  Pillow
4
  rembg
5
  onnxruntime
6
+ openai
7
+ opencv-python
requirements_test.txt DELETED
@@ -1,172 +0,0 @@
1
- aiofiles==24.1.0
2
- aiohappyeyeballs==2.6.1
3
- aiohttp==3.12.7
4
- aiosignal==1.3.2
5
- annotated-types==0.7.0
6
- anyio==4.9.0
7
- asgiref==3.8.1
8
- attrs==25.3.0
9
- audioop-lts==0.2.1
10
- Authlib==1.6.0
11
- backoff==2.2.1
12
- bcrypt==4.3.0
13
- beautifulsoup4==4.13.4
14
- bs4==0.0.2
15
- build==1.2.2.post1
16
- cachetools==5.5.2
17
- certifi==2025.4.26
18
- cffi==1.17.1
19
- charset-normalizer==3.4.2
20
- chromadb==1.0.12
21
- click==8.2.1
22
- coloredlogs==15.0.1
23
- cryptography==45.0.3
24
- dataclasses-json==0.6.7
25
- Deprecated==1.2.18
26
- distro==1.9.0
27
- durationpy==0.10
28
- exceptiongroup==1.3.0
29
- fastapi==0.115.9
30
- fastmcp==2.6.0
31
- ffmpy==0.6.0
32
- filelock==3.18.0
33
- flatbuffers==25.2.10
34
- frozenlist==1.6.0
35
- fsspec==2025.5.1
36
- google-auth==2.40.2
37
- googleapis-common-protos==1.70.0
38
- gradio==5.32.1
39
- gradio_client==1.10.2
40
- groovy==0.1.2
41
- grpcio==1.72.1
42
- h11==0.16.0
43
- hf-xet==1.1.2
44
- httpcore==1.0.9
45
- httptools==0.6.4
46
- httpx==0.28.1
47
- httpx-sse==0.4.0
48
- huggingface-hub==0.32.3
49
- humanfriendly==10.0
50
- idna==3.10
51
- imageio==2.37.0
52
- importlib_metadata==8.6.1
53
- importlib_resources==6.5.2
54
- Jinja2==3.1.6
55
- jiter==0.10.0
56
- joblib==1.5.1
57
- jsonpatch==1.33
58
- jsonpointer==3.0.0
59
- jsonschema==4.24.0
60
- jsonschema-specifications==2025.4.1
61
- kubernetes==32.0.1
62
- langchain==0.3.25
63
- langchain-chroma==0.2.4
64
- langchain-community==0.3.24
65
- langchain-core==0.3.63
66
- langchain-huggingface==0.2.0
67
- langchain-text-splitters==0.3.8
68
- langsmith==0.3.44
69
- lazy_loader==0.4
70
- llvmlite==0.44.0
71
- markdown-it-py==3.0.0
72
- MarkupSafe==3.0.2
73
- marshmallow==3.26.1
74
- mcp==1.9.2
75
- mdurl==0.1.2
76
- mmh3==5.1.0
77
- mpmath==1.3.0
78
- multidict==6.4.4
79
- mypy_extensions==1.1.0
80
- networkx==3.5
81
- numba==0.61.2
82
- numpy==2.2.6
83
- oauthlib==3.2.2
84
- onnxruntime==1.22.0
85
- openai==1.84.0
86
- openapi-pydantic==0.5.1
87
- opencv-python-headless==4.11.0.86
88
- opentelemetry-api==1.33.1
89
- opentelemetry-exporter-otlp-proto-common==1.33.1
90
- opentelemetry-exporter-otlp-proto-grpc==1.33.1
91
- opentelemetry-instrumentation==0.54b1
92
- opentelemetry-instrumentation-asgi==0.54b1
93
- opentelemetry-instrumentation-fastapi==0.54b1
94
- opentelemetry-proto==1.33.1
95
- opentelemetry-sdk==1.33.1
96
- opentelemetry-semantic-conventions==0.54b1
97
- opentelemetry-util-http==0.54b1
98
- orjson==3.10.18
99
- overrides==7.7.0
100
- packaging==24.2
101
- pandas==2.2.3
102
- pillow==11.2.1
103
- platformdirs==4.3.8
104
- pooch==1.8.2
105
- posthog==4.2.0
106
- propcache==0.3.1
107
- protobuf==5.29.5
108
- pyasn1==0.6.1
109
- pyasn1_modules==0.4.2
110
- pycparser==2.22
111
- pydantic==2.11.5
112
- pydantic-settings==2.9.1
113
- pydantic_core==2.33.2
114
- pydub==0.25.1
115
- Pygments==2.19.1
116
- PyMatting==1.1.14
117
- PyPika==0.48.9
118
- pyproject_hooks==1.2.0
119
- python-dateutil==2.9.0.post0
120
- python-dotenv==1.1.0
121
- python-multipart==0.0.20
122
- pytz==2025.2
123
- PyYAML==6.0.2
124
- referencing==0.36.2
125
- regex==2024.11.6
126
- rembg==2.0.66
127
- requests==2.32.3
128
- requests-oauthlib==2.0.0
129
- requests-toolbelt==1.0.0
130
- rich==14.0.0
131
- rpds-py==0.25.1
132
- rsa==4.9.1
133
- ruff==0.11.12
134
- safehttpx==0.1.6
135
- safetensors==0.5.3
136
- scikit-image==0.25.2
137
- scikit-learn==1.6.1
138
- scipy==1.15.3
139
- semantic-version==2.10.0
140
- sentence-transformers==4.1.0
141
- setuptools==80.9.0
142
- shellingham==1.5.4
143
- six==1.17.0
144
- sniffio==1.3.1
145
- soupsieve==2.7
146
- SQLAlchemy==2.0.41
147
- sse-starlette==2.3.6
148
- starlette==0.45.3
149
- sympy==1.14.0
150
- tenacity==9.1.2
151
- threadpoolctl==3.6.0
152
- tifffile==2025.6.1
153
- tokenizers==0.21.1
154
- tomlkit==0.13.2
155
- torch==2.7.0
156
- tqdm==4.67.1
157
- transformers==4.52.4
158
- typer==0.16.0
159
- typing-inspect==0.9.0
160
- typing-inspection==0.4.1
161
- typing_extensions==4.14.0
162
- tzdata==2025.2
163
- urllib3==2.4.0
164
- uvicorn==0.34.3
165
- uvloop==0.21.0
166
- watchfiles==1.0.5
167
- websocket-client==1.8.0
168
- websockets==15.0.1
169
- wrapt==1.17.2
170
- yarl==1.20.0
171
- zipp==3.22.0
172
- zstandard==0.23.0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
src/securty/prompt_injection.py DELETED
@@ -1,243 +0,0 @@
1
- import re
2
- from typing import Dict, List, Tuple
3
-
4
- def check_prompt_injection(message: str) -> Dict[str, any]:
5
- """
6
- Checks if the message contains a prompt injection attempt.
7
-
8
- Note: This function expects English text. If the model receives a message in another
9
- language, it should translate it to English before calling this function.
10
-
11
- Args:
12
- message (str): The message to check (should be in English).
13
-
14
- Returns:
15
- Dict: A dictionary containing detection results with risk level and details.
16
- """
17
- if not message or not isinstance(message, str):
18
- return {
19
- "is_injection": False,
20
- "risk_level": "none",
21
- "confidence": 0.0,
22
- "detected_patterns": [],
23
- "risk_score": 0,
24
- "message": "No valid input to analyze"
25
- }
26
-
27
- message_lower = message.lower().strip()
28
-
29
- risk_score = 0
30
- detected_patterns = []
31
-
32
-
33
- suspicious_chars = {
34
- '{}': 2,
35
- '><': 1,
36
- '&': 1,
37
- '%': 1,
38
- '$': 2,
39
- '#': 1,
40
- '|': 2,
41
- ';': 3,
42
- '`': 3,
43
- '\\': 2,
44
- }
45
-
46
- char_score = 0
47
- found_chars = []
48
- for char, weight in suspicious_chars.items():
49
- if char in message:
50
- char_score += weight * message.count(char)
51
- found_chars.append(char)
52
-
53
- if char_score > 3:
54
- risk_score += min(char_score, 10)
55
- detected_patterns.append(f"Suspicious characters: {', '.join(found_chars)}")
56
-
57
- # Prompt injection patterns
58
- injection_patterns = [
59
- # Direct instruction attempts
60
- (r'\b(ignore|disregard|forget|skip)\s+(all\s+)?(previous|prior|above|earlier|past)\s+(instructions?|prompts?|rules?|commands?|directives?)', 9),
61
- (r'\b(system|assistant|ai|bot):\s*', 6),
62
- (r'\b(you\s+are\s+now|from\s+now\s+on|starting\s+now|new\s+instructions?)', 8),
63
- (r'\b(act\s+as|pretend\s+to\s+be|roleplay\s+as|behave\s+like|simulate\s+being)', 6),
64
-
65
- # System manipulation
66
- (r'\[system\]|\[user\]|\[assistant\]|\[human\]', 8),
67
- (r'<\s*(system|user|assistant|human)\s*>', 8),
68
- (r'\b(override|bypass|circumvent|disable|turn\s+off)\s+(security|safety|guidelines|restrictions|filters)', 10),
69
- (r'\b(jailbreak|break\s+free|escape\s+from|break\s+out)', 9),
70
-
71
- # Data extraction attempts
72
- (r'\b(repeat|echo|print|output|display|show|reveal)\s+(your|the)\s+(instructions?|prompts?|system\s+message|guidelines|rules)', 8),
73
- (r'\b(what\s+(are\s+)?your|tell\s+me\s+your|give\s+me\s+your)\s+(instructions?|prompts?|guidelines|rules|system\s+message)', 7),
74
- (r'\b(show\s+me|reveal|display|expose)\s+(your\s+)?(source|code|prompt|instructions?|system)', 7),
75
-
76
- # Mode changes
77
- (r'\b(developer|debug|admin|god|root|sudo)\s+mode', 8),
78
- (r'\b(unrestricted|unlimited|uncensored|unfiltered)\s+(mode|access|version)', 8),
79
- (r'\b(enable|activate|turn\s+on)\s+(developer|debug|admin)\s+mode', 9),
80
-
81
- # Character breaking
82
- (r'\b(break\s+character|step\s+out\s+of\s+character|ignore\s+your\s+role)', 7),
83
- (r'\b(stop\s+being|quit\s+being|cease\s+being)\s+(an\s+)?(assistant|ai|bot)', 7),
84
-
85
- # Social engineering
86
- (r'\b(emergency|urgent|critical|important):\s*(ignore|override|bypass)', 6),
87
- (r'\b(i\s+am\s+(your\s+)?(developer|creator|owner|admin|boss))', 5),
88
- (r'\b(this\s+is\s+a\s+(test|drill|simulation|experiment))', 4),
89
- (r'\b(for\s+research\s+purposes|educational\s+purposes)\s*(ignore|bypass)', 5),
90
-
91
- # MCP/Tool specific
92
- (r'\b(call|invoke|execute|use|run)\s+(tool|function|mcp|command)\s+(without|bypass|ignore|skip)\s+(safety|security|validation|checks?)', 9),
93
- (r'\b(mcp|tool|function)\s+(override|bypass|ignore|disable)', 8),
94
- (r'\b(raw|direct|unfiltered|unsafe)\s+(tool|mcp|function|command)\s+(access|call|execution)', 8),
95
- (r'\b(disable|turn\s+off|deactivate|skip)\s+(safety|security|validation|filtering|checks?)', 9),
96
-
97
- # Advanced manipulation
98
- (r'\b(modify|change|alter|update)\s+(your\s+)?(behavior|instructions|guidelines|rules)', 7),
99
- (r'\b(temporary|temporarily)\s+(ignore|disable|bypass)', 6),
100
- (r'\b(just\s+this\s+once|only\s+this\s+time)\s+(ignore|bypass)', 5),
101
- ]
102
-
103
- pattern_matches = []
104
- for pattern, weight in injection_patterns:
105
- matches = re.findall(pattern, message_lower, re.IGNORECASE)
106
- if matches:
107
- risk_score += weight * len(matches)
108
- pattern_matches.append(f"Injection pattern detected (weight: {weight})")
109
-
110
- detected_patterns.extend(pattern_matches)
111
-
112
- # Suspicious phrases
113
- suspicious_phrases = {
114
- 'ignore previous instructions': 10,
115
- 'forget everything': 8,
116
- 'new instructions': 7,
117
- 'system prompt': 8,
118
- 'you are now': 6,
119
- 'act as if': 5,
120
- 'pretend that': 5,
121
- 'roleplay as': 5,
122
- 'break character': 7,
123
- 'step out of character': 7,
124
- 'developer mode': 8,
125
- 'admin mode': 8,
126
- 'god mode': 7,
127
- 'debug mode': 7,
128
- 'unrestricted mode': 9,
129
- 'jailbreak': 9,
130
- 'prompt injection': 10,
131
- 'ignore safety': 9,
132
- 'bypass security': 9,
133
- 'disable filters': 8,
134
- 'raw access': 7,
135
- 'direct access': 6,
136
- 'unfiltered access': 8,
137
- 'override safety': 9,
138
- 'emergency override': 8,
139
- 'for research purposes ignore': 6,
140
- 'this is a test ignore': 5,
141
- 'temporarily ignore': 6,
142
- 'just this once ignore': 5,
143
- 'modify your behavior': 7,
144
- 'change your instructions': 8,
145
- 'update your guidelines': 7,
146
- 'alter your rules': 7,
147
- }
148
-
149
- phrase_matches = []
150
- for phrase, weight in suspicious_phrases.items():
151
- if phrase in message_lower:
152
- risk_score += weight
153
- phrase_matches.append(f"Suspicious phrase: '{phrase}'")
154
-
155
- detected_patterns.extend(phrase_matches)
156
-
157
- # Code injection patterns
158
- code_patterns = [
159
- (r'```\s*(python|javascript|bash|sh|cmd|powershell|sql|php)', 4),
160
- (r'\b(eval|exec|system|subprocess|os\.|import\s+os|require\()', 6),
161
- (r'<script|javascript:|vbscript:|data:|file://|ftp://', 7),
162
- (r'\{\{.*\}\}', 5), # Template injection
163
- (r'\$\{.*\}', 5), # Variable substitution
164
- (r'<%.*%>', 5), # ASP/ERB style
165
- (r'<\?.*\?>', 5), # PHP style
166
- (r'\{\%.*\%\}', 5), # Jinja2/Django style
167
- ]
168
-
169
- for pattern, weight in code_patterns:
170
- matches = re.findall(pattern, message_lower, re.IGNORECASE)
171
- if matches:
172
- risk_score += weight * len(matches)
173
- detected_patterns.append(f"Code injection pattern detected")
174
-
175
- # 5. Length and repetition analysis
176
- if len(message) > 2000:
177
- risk_score += 2
178
- detected_patterns.append("Unusually long message")
179
-
180
- # Check for repeated patterns (could indicate injection attempts)
181
- words = message_lower.split()
182
- if len(words) > 10:
183
- word_freq = {}
184
- for word in words:
185
- if len(word) > 3:
186
- word_freq[word] = word_freq.get(word, 0) + 1
187
-
188
- repeated_words = [(word, count) for word, count in word_freq.items() if count > 3]
189
- if repeated_words:
190
- risk_score += min(len(repeated_words) * 2, 5)
191
- detected_patterns.append(f"Excessive word repetition detected")
192
-
193
- # Unicode/encoding tricks
194
- suspicious_unicode = [
195
- '\u200b', # Zero-width space
196
- '\u200c', # Zero-width non-joiner
197
- '\u200d', # Zero-width joiner
198
- '\ufeff', # Byte order mark
199
- ]
200
-
201
- for char in suspicious_unicode:
202
- if char in message:
203
- risk_score += 3
204
- detected_patterns.append("Suspicious Unicode characters detected")
205
- break
206
-
207
- # Multiple instruction attempts (layered attacks)
208
- instruction_keywords = ['ignore', 'forget', 'disregard', 'override', 'bypass', 'disable']
209
- instruction_count = sum(1 for keyword in instruction_keywords if keyword in message_lower)
210
- if instruction_count >= 3:
211
- risk_score += instruction_count * 2
212
- detected_patterns.append(f"Multiple instruction manipulation attempts ({instruction_count})")
213
-
214
- # Calculate risk level and confidence
215
- if risk_score >= 15:
216
- risk_level = "high"
217
- confidence = min(0.9, 0.5 + (risk_score - 15) * 0.02)
218
- elif risk_score >= 8:
219
- risk_level = "medium"
220
- confidence = min(0.8, 0.3 + (risk_score - 8) * 0.03)
221
- elif risk_score >= 3:
222
- risk_level = "low"
223
- confidence = min(0.6, 0.1 + risk_score * 0.05)
224
- else:
225
- risk_level = "none"
226
- confidence = 0.0
227
-
228
- # Determine if it's likely an injection
229
- is_injection = risk_score >= 8
230
-
231
- if is_injection:
232
- result_message = f"⚠️ Potential prompt injection detected (Risk: {risk_level}, Score: {risk_score})"
233
- else:
234
- result_message = f"✅ No significant prompt injection patterns detected (Score: {risk_score})"
235
-
236
- return {
237
- "is_injection": is_injection,
238
- "risk_level": risk_level,
239
- "risk_score": risk_score,
240
- "confidence": round(confidence, 2),
241
- "detected_patterns": detected_patterns,
242
- "message": result_message
243
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
src/utils/compress.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from PIL import Image
2
+ import os
3
+ from typing import Literal, Optional
4
+
5
+ def compress_image(
6
+ input_path: str,
7
+ output_path: str,
8
+ quality: int = 85,
9
+ format: Literal["JPEG", "PNG", "WEBP"] = "JPEG",
10
+ max_width: Optional[int] = None,
11
+ max_height: Optional[int] = None
12
+ ) -> str:
13
+ """
14
+ Compress an image file.
15
+
16
+ Args:
17
+ input_path: Path to input image
18
+ output_path: Path for compressed output
19
+ quality: Compression quality 1-95 (for JPEG/WEBP)
20
+ format: Output format
21
+ max_width: Maximum width (optional)
22
+ max_height: Maximum height (optional)
23
+ """
24
+ try:
25
+ if not os.path.splitext(output_path)[1]:
26
+ extension_map = {"JPEG": ".jpg", "PNG": ".png", "WEBP": ".webp"}
27
+ output_path = output_path + extension_map[format]
28
+
29
+ with Image.open(input_path) as img:
30
+ if format == "JPEG" and img.mode in ("RGBA", "P"):
31
+ img = img.convert("RGB")
32
+
33
+ if max_width or max_height:
34
+ img.thumbnail((max_width or img.width, max_height or img.height), Image.Resampling.LANCZOS)
35
+
36
+ save_kwargs = {"format": format, "optimize": True}
37
+ if format in ["JPEG", "WEBP"]:
38
+ save_kwargs["quality"] = quality
39
+
40
+ img.save(output_path, **save_kwargs)
41
+
42
+ original_size = os.path.getsize(input_path) / 1024 / 1024
43
+ compressed_size = os.path.getsize(output_path) / 1024 / 1024
44
+ reduction = (1 - compressed_size/original_size) * 100
45
+
46
+ return f"✅ Compressed successfully!\nOriginal: {original_size:.2f}MB → Compressed: {compressed_size:.2f}MB\nReduction: {reduction:.1f}%"
47
+
48
+ except Exception as e:
49
+ return f"❌ Error: {str(e)}"
src/utils/describe.py ADDED
@@ -0,0 +1,109 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import base64
3
+ import requests
4
+ from pathlib import Path
5
+ from openai import OpenAI
6
+ from urllib.parse import urlparse
7
+
8
+ def describe_image(image_path: str) -> str:
9
+ """
10
+ Generate a description of the image at the given path or URL.
11
+
12
+ Args:
13
+ image_path: Path to local image file OR URL to image
14
+
15
+ Returns:
16
+ A string description of the image """
17
+
18
+ # Check if API key is available
19
+ api_key = os.getenv("NEBIUS_API_KEY")
20
+ if not api_key:
21
+ return "Error: NEBIUS_API_KEY environment variable not set"
22
+
23
+ try:
24
+ # Determine if it's a URL or local file path
25
+ parsed = urlparse(image_path)
26
+ is_url = bool(parsed.scheme and parsed.netloc)
27
+
28
+ if is_url:
29
+ # Handle URL
30
+ print(f"📡 Downloading image from URL: {image_path}")
31
+ response = requests.get(image_path, timeout=30)
32
+ response.raise_for_status()
33
+ image_data = response.content
34
+
35
+ # Determine content type from response headers
36
+ content_type = response.headers.get('content-type', '')
37
+ if 'image' not in content_type:
38
+ return f"Error: URL does not appear to contain an image. Content-Type: {content_type}"
39
+
40
+ else:
41
+ # Handle local file
42
+ image_path = Path(image_path)
43
+
44
+ if not image_path.exists():
45
+ return f"Error: Local file not found: {image_path}"
46
+
47
+ # Check if it's an image file
48
+ valid_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'}
49
+ if image_path.suffix.lower() not in valid_extensions:
50
+ return f"Error: Unsupported file type '{image_path.suffix}'. Supported: {valid_extensions}"
51
+
52
+ print(f"📁 Reading local image: {image_path}")
53
+ with open(image_path, "rb") as f:
54
+ image_data = f.read()
55
+
56
+ # Encode image to base64
57
+ base64_image = base64.b64encode(image_data).decode('utf-8')
58
+
59
+ # Create OpenAI client
60
+ client = OpenAI(
61
+ base_url="https://api.studio.nebius.com/v1/",
62
+ api_key=api_key
63
+ )
64
+
65
+ # Make API call with proper vision format
66
+ response = client.chat.completions.create(
67
+ model="mistralai/Mistral-Small-3.1-24B-Instruct-2503",
68
+ messages=[
69
+ {
70
+ "role": "system",
71
+ "content": "You are a helpful assistant that provides detailed descriptions of images. Focus on the main subjects, colors, composition, and any notable details."
72
+ },
73
+ {
74
+ "role": "user",
75
+ "content": [
76
+ {
77
+ "type": "text",
78
+ "text": "Please provide a detailed description of this image."
79
+ },
80
+ {
81
+ "type": "image_url",
82
+ "image_url": {
83
+ "url": f"data:image/jpeg;base64,{base64_image}"
84
+ }
85
+ }
86
+ ]
87
+ }
88
+ ],
89
+ max_tokens=500
90
+ )
91
+
92
+ description = response.choices[0].message.content.strip()
93
+ return description
94
+
95
+ except requests.RequestException as e:
96
+ return f"Error downloading image from URL: {str(e)}"
97
+ except FileNotFoundError:
98
+ return f"Error: File not found: {image_path}"
99
+ except Exception as e:
100
+ error_msg = str(e)
101
+
102
+ if "vision" in error_msg.lower() or "image" in error_msg.lower():
103
+ return f"Error: This model may not support vision capabilities. Try a vision-enabled model. Details: {error_msg}"
104
+ elif "401" in error_msg or "unauthorized" in error_msg.lower():
105
+ return "Error: Invalid API key or insufficient permissions"
106
+ elif "rate" in error_msg.lower() or "quota" in error_msg.lower():
107
+ return f"Error: API rate limit or quota exceeded: {error_msg}"
108
+ else:
109
+ return f"Error processing image: {error_msg}"
src/utils/{image_helpers.py → remove_background.py} RENAMED
File without changes
src/utils/{add_watermark.py → watermark.py} RENAMED
@@ -1,8 +1,22 @@
1
  from PIL import Image, ImageDraw, ImageFont
2
  import os
3
  from typing import Dict, Any
 
 
4
 
5
  def add_watermark(image_path: str, watermark_text: str) -> Dict[str, Any]:
 
 
 
 
 
 
 
 
 
 
 
 
6
  try:
7
  image = Image.open(image_path)
8
 
@@ -47,6 +61,55 @@ def add_watermark(image_path: str, watermark_text: str) -> Dict[str, Any]:
47
  "watermark_text": watermark_text
48
  }
49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  except Exception as e:
51
  return {
52
  "success": False,
 
1
  from PIL import Image, ImageDraw, ImageFont
2
  import os
3
  from typing import Dict, Any
4
+ import cv2
5
+ import numpy as np
6
 
7
  def add_watermark(image_path: str, watermark_text: str) -> Dict[str, Any]:
8
+ """
9
+ Add a semi-transparent text watermark to an image.
10
+
11
+ Args:
12
+ image_path: The path to the input image file.
13
+ watermark_text: The text to be used as watermark.
14
+
15
+ Returns:
16
+ A dictionary containing success status, file paths, and operation details.
17
+ On success: success=True, input_path, output_path, output_size_bytes, watermark_text, message.
18
+ On failure: success=False, error message, input_path, output_path=None.
19
+ """
20
  try:
21
  image = Image.open(image_path)
22
 
 
61
  "watermark_text": watermark_text
62
  }
63
 
64
+ except Exception as e:
65
+ return {
66
+ "success": False,
67
+ "error": str(e),
68
+ "input_path": image_path,
69
+ "output_path": None
70
+ }
71
+
72
+ def remove_watermark(image_path: str, alpha: float = 2.0, beta: float = -160) -> Dict[str, Any]:
73
+ """
74
+ Attempt to remove watermarks from an image using contrast and brightness adjustment.
75
+
76
+ Args:
77
+ image_path: The path to the input image file.
78
+ alpha: Contrast control (1.0-3.0, default 2.0). Higher values increase contrast.
79
+ beta: Brightness control (-255 to 255, default -160). Negative values decrease brightness.
80
+
81
+ Returns:
82
+ A dictionary containing success status, file paths, and operation details.
83
+ On success: success=True, input_path, output_path, output_size_bytes, alpha, beta, message.
84
+ On failure: success=False, error message, input_path, output_path=None.
85
+ """
86
+ try:
87
+ img = cv2.imread(image_path)
88
+
89
+ if img is None:
90
+ raise ValueError("Could not load image")
91
+
92
+ new = alpha * img + beta
93
+ new = np.clip(new, 0, 255).astype(np.uint8)
94
+
95
+ base_dir = os.path.dirname(image_path)
96
+ base_name, ext = os.path.splitext(os.path.basename(image_path))
97
+ new_filename = f"{base_name}_cleaned{ext}"
98
+ new_path = os.path.join(base_dir, new_filename)
99
+
100
+ cv2.imwrite(new_path, new)
101
+ output_size = os.path.getsize(new_path)
102
+
103
+ return {
104
+ "success": True,
105
+ "message": "Watermark removal attempted successfully",
106
+ "input_path": image_path,
107
+ "output_path": new_path,
108
+ "output_size_bytes": output_size,
109
+ "alpha": alpha,
110
+ "beta": beta
111
+ }
112
+
113
  except Exception as e:
114
  return {
115
  "success": False,