Fix minor bugs

README.md

@@ -7,7 +7,7 @@ This repository contains a Prodigy plugin for recipes that involve Metas [Segmen
 You can install this plugin via `pip`. 
 
 ```
-pip install "prodigy-ann @ git+https://github.com/explosion/prodigy-segment"
+pip install "prodigy-segment @ git+https://github.com/explosion/prodigy-segment"
 ```
 
 To learn more about this plugin, you can check the [Prodigy docs](https://prodi.gy/docs/plugins/#segment).

prodigy_segment/__init__.py

@@ -14,7 +14,7 @@
 from prodigy.types import LabelsType, SourceType, TaskType
 from prodigy.util import log, msg
 from prodigy_segment.segment_anything import sam_model_registry, SamPredictor
-
+from typing import Dict, Union
 
 HTML = """
 <link
@@ -125,12 +125,13 @@ def calculate_masks(box_coordinates: List, predictor: SamPredictor, pil_image: I
     return masks
 
 
-def get_base64_string(img_str: str):
+def get_base64_string(img: Union[Dict[str, str], str]):
     # This looks hacky at first glance, but the reasoning here is that the schema
     # per https://en.wikipedia.org/wiki/Data_URI_scheme#Syntax looks like this:
     # data:[<media type>][;charset=<character set>][;base64],<data>
     # The encoding will always end with base64, so that's the easy place to cut. 
     # Otherwise we risk assuming a media type or characterset.
+    img_str = img["image"] if isinstance(img, dict) else img
     str_idx = img_str.find("base64,") + 7
     return img_str[str_idx:]
 

setup.cfg

@@ -21,8 +21,7 @@ install_requires =
 [options.entry_points]
 prodigy_recipes =
     segment.image.manual = prodigy_segment:segment_image_manual
-    segment.cache = prodigy_segment:segment_cache
-
+    segment.cache = prodigy_segment:segment_fill_cache
 [bdist_wheel]
 universal = true