add remove_bg function

agent_jarvis/TestPage/index.html

@@ -13,7 +13,7 @@
       <div id="status"></div>
       <div id="bottom_bar">
         <div id="form">
-          <input id="input" autocomplete="off" />
+          <input id="input" autocomplete="off" placeholder="Drag an image here to send image" />
           <button id="send_button">Send</button>
           <button id="reset_button">ClearMemory</button>
           <button id="test_button">AutoTest</button>

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 = `<img src="${file}" alt="Image" class="replied-img">`;
+          item.classList.add("align_right");
+          messages.appendChild(item);
+          scrollElementToEnd(messages);
+        }
+      };
+      reader.readAsDataURL(file);
+    }
+  });
 })();

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

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")
 

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")

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:

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

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()

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"