Fix Issues (#140)

* convert to rbg for jpg

* handle heic images

* ignore import type error

* add as default import for remote code exec

pyproject.toml

@@ -37,6 +37,7 @@ ipykernel = "^6.29.4"
 e2b = "^0.17.1"
 e2b-code-interpreter = "^0.0.9"
 tenacity = "^8.3.0"
+pillow-heif = "^0.16.0"
 
 [tool.poetry.group.dev.dependencies]
 autoflake = "1.*"

vision_agent/agent/vision_agent.py

@@ -7,8 +7,8 @@
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Sequence, Union, cast
 
-from PIL import Image
 from langsmith import traceable
+from PIL import Image
 from rich.console import Console
 from rich.style import Style
 from rich.syntax import Syntax
@@ -43,6 +43,7 @@ class DefaultImports:
 
  common_imports = [
  "from typing import *",
+ "register_heif_opener()",
  ]
 
  @staticmethod

vision_agent/tools/__init__.py

@@ -17,6 +17,8 @@
  grounding_dino,
  grounding_sam,
  load_image,
+ loca_visual_prompt_counting,
+ loca_zero_shot_counting,
  ocr,
  overlay_bounding_boxes,
  overlay_heat_map,
@@ -25,8 +27,6 @@
  save_image,
  save_json,
  save_video,
- loca_visual_prompt_counting,
- loca_zero_shot_counting,
  vit_image_classification,
  vit_nsfw_classification,
 )

vision_agent/tools/tools.py

@@ -13,6 +13,7 @@
 import requests
 from moviepy.editor import ImageSequenceClip
 from PIL import Image, ImageDraw, ImageFont
+from pillow_heif import register_heif_opener # type: ignore
 
 from vision_agent.tools.tool_utils import _send_inference_request
 from vision_agent.utils import extract_frames_from_video
@@ -26,6 +27,8 @@
  rle_decode,
 )
 
+register_heif_opener()
+
 COLORS = [
  (158, 218, 229),
  (219, 219, 141),
@@ -659,7 +662,7 @@ def save_image(image: np.ndarray, file_path: str) -> None:
  """
  from IPython.display import display
 
- pil_image = Image.fromarray(image.astype(np.uint8))
+ pil_image = Image.fromarray(image.astype(np.uint8)).convert("RGB")
  display(pil_image)
  pil_image.save(file_path)