-
+
diff --git a/agent_jarvis/TestPage/js/main.js b/agent_jarvis/TestPage/js/main.js
index 0bd4144..6bede6c 100644
--- a/agent_jarvis/TestPage/js/main.js
+++ b/agent_jarvis/TestPage/js/main.js
@@ -187,5 +187,36 @@
if (e.key === 'Enter') {
submitMessage();
}
- })
+ });
+
+ input.addEventListener('drop', function(e) {
+ e.preventDefault();
+ e.stopPropagation();
+
+ const files = e.dataTransfer.files;
+ for (let i = 0; i < files.length; i++) {
+ const file = files[i];
+ const reader = new FileReader();
+ reader.onload = function(e) {
+ const file = e.currentTarget.result;
+ if (file.startsWith('data:image')) {
+ const base64Index = file.indexOf('base64,');
+ const base64Content = file.substring(base64Index + 7)
+ // Send to jarvis
+ socket.emit(
+ "chat_message",
+ makeChatMessage(base64Content, 'image')
+ );
+
+ // display on the page
+ var item = document.createElement("li");
+ item.innerHTML = `
`;
+ item.classList.add("align_right");
+ messages.appendChild(item);
+ scrollElementToEnd(messages);
+ }
+ };
+ reader.readAsDataURL(file);
+ }
+ });
})();
diff --git a/agent_jarvis/dockerfile b/agent_jarvis/dockerfile
index 6afbd68..f647592 100644
--- a/agent_jarvis/dockerfile
+++ b/agent_jarvis/dockerfile
@@ -31,6 +31,7 @@ RUN echo "openai==0.27.6" >> requirements.txt \
&& echo "llama-index==0.6.7" >> requirements.txt \
&& echo "youtube-transcript-api==0.6.0" >> requirements.txt \
&& echo "tweepy==4.14.0" >> requirements.txt \
+ && echo "rembg==2.0.45" >> requirements.txt \
&& echo "google-api-python-client==2.86.0" >> requirements.txt \
&& echo "google-auth-httplib2==0.1.0" >> requirements.txt \
&& echo "google-auth-oauthlib==1.0.0" >> requirements.txt
diff --git a/agent_jarvis/jarvis/functional_modules/caller_context.py b/agent_jarvis/jarvis/functional_modules/caller_context.py
index 24ca273..24ada24 100644
--- a/agent_jarvis/jarvis/functional_modules/caller_context.py
+++ b/agent_jarvis/jarvis/functional_modules/caller_context.py
@@ -18,6 +18,12 @@ class CallerContext:
return f"{of}"
return ""
+ def get_last_image(self) -> str:
+ raise NotImplementedError("Function not implemented")
+
+ def set_last_image(self, img: str):
+ raise NotImplementedError("Function not implemented")
+
async def reply_text(self, msg):
raise NotImplementedError("Function not implemented")
diff --git a/agent_jarvis/jarvis/gateway/session.py b/agent_jarvis/jarvis/gateway/session.py
index 51f13ad..eb5d16e 100644
--- a/agent_jarvis/jarvis/gateway/session.py
+++ b/agent_jarvis/jarvis/gateway/session.py
@@ -65,6 +65,7 @@ class Session(CallerContext):
_sio: SioConnection = None
_session_id: str = None
_tz_offset: int = 0 # timezone offset, in hours. = local_time - UTC0
+ _last_image: str = None
_history_dir: str = None
@@ -163,6 +164,12 @@ class Session(CallerContext):
def get_tz_offset(self):
return self._tz_offset
+ def get_last_image(self) -> str:
+ return self._last_image
+
+ def set_last_image(self, img: str):
+ self._last_image = img
+
async def reply_text(self, msg):
if self._message_id_ is None:
raise function_error.FunctionError(function_error.EC_RESET, "Reset")
diff --git a/agent_jarvis/jarvis/main.py b/agent_jarvis/jarvis/main.py
index 98616b3..d2a9771 100644
--- a/agent_jarvis/jarvis/main.py
+++ b/agent_jarvis/jarvis/main.py
@@ -109,6 +109,8 @@ def run_server_mode():
session.set_tz_offset(offset)
elif msg.message_type == 'text':
await session.on_chat_message(msg)
+ elif msg.message_type == 'image':
+ session.set_last_image(msg.message_content)
app.router.add_static('/js', './TestPage/js')
app.router.add_static('/css', './TestPage/css')
@@ -156,6 +158,8 @@ async def run_client_mode(session_map: dict[str, Session]):
session.set_tz_offset(offset)
elif msg.message_type == 'text':
await session.on_chat_message(msg)
+ elif msg.message_type == 'image':
+ session.set_last_image(msg.message_content)
await sio.connect(CFG.bot_server_url)
try:
diff --git a/agent_jarvis/requirements.txt b/agent_jarvis/requirements.txt
index 55d0571..c4dd225 100644
--- a/agent_jarvis/requirements.txt
+++ b/agent_jarvis/requirements.txt
@@ -9,6 +9,7 @@ python-dotenv==1.0.0
llama-index==0.6.7
youtube-transcript-api==0.6.0
tweepy==4.14.0
+rembg==2.0.45
google-api-python-client==2.86.0
google-auth-httplib2==0.1.0
diff --git a/example_modules/demo_modules/rembg.module.py b/example_modules/demo_modules/rembg.module.py
new file mode 100644
index 0000000..95bf95c
--- /dev/null
+++ b/example_modules/demo_modules/rembg.module.py
@@ -0,0 +1,41 @@
+import io
+import base64
+from typing import List
+
+from jarvis.functional_modules.caller_context import CallerContext
+from jarvis.functional_modules.functional_module import functional_module
+from jarvis.logger import logger
+
+
+def reg_or_not():
+ try:
+ from rembg import new_session, remove
+ from PIL import Image
+ except ImportError as e:
+ logger.warn(f"rembg or PIL not installed, remove_bg will not be available")
+ return
+
+ session = new_session(model_name="u2net")
+
+ @functional_module(
+ name="remove_bg",
+ description="Remove the background of last image",
+ signature={})
+ async def remove_bg(context: CallerContext):
+ img_base64 = context.get_last_image()
+ if img_base64 is None:
+ await context.reply_text("You need to give me an image first")
+ return "Failed"
+ bytes_io = io.BytesIO(base64.b64decode(img_base64))
+ img = Image.open(bytes_io)
+ output = remove(img, session=session)
+ output_bytesio = io.BytesIO()
+ output.save(output_bytesio, 'PNG')
+ output_bytes = output_bytesio.getvalue()
+ result = base64.b64encode(output_bytes).decode()
+
+ context.set_last_image(result)
+ await context.reply_image_base64(result)
+ return "Success"
+
+reg_or_not()
diff --git a/example_modules/demo_modules/stable_diffusion.module.py b/example_modules/demo_modules/stable_diffusion.module.py
index 78c3e08..e22dc0a 100644
--- a/example_modules/demo_modules/stable_diffusion.module.py
+++ b/example_modules/demo_modules/stable_diffusion.module.py
@@ -260,6 +260,7 @@ NOTE: Just reply using these information, don't ask me anything.
logger.debug("Start calling stable_diffusion")
img = await call_sd(sd_params)
logger.debug("End calling stable_diffusion")
+ context.set_last_image(img)
await context.reply_image_base64(img)
return "Success"