clarifai 9.10.1__py3-none-any.whl → 9.10.3__py3-none-any.whl

clarifai/client/user.py

@@ -15,18 +15,23 @@ from clarifai.utils.logging import get_logger
 class User(Lister, BaseClient):
   """User is a class that provides access to Clarifai API endpoints related to user information."""
 
-  def __init__(self, user_id: str = "", base_url: str = "https://api.clarifai.com", **kwargs):
+  def __init__(self,
+               user_id: str = None,
+               base_url: str = "https://api.clarifai.com",
+               pat: str = None,
+               **kwargs):
     """Initializes an User object.
 
     Args:
         user_id (str): The user ID for the user to interact with.
         base_url (str): Base API url. Default "https://api.clarifai.com"
+        pat (str): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
         **kwargs: Additional keyword arguments to be passed to the User.
     """
     self.kwargs = {**kwargs, 'id': user_id}
     self.user_info = resources_pb2.User(**self.kwargs)
     self.logger = get_logger(logger_level="INFO", name=__name__)
-    BaseClient.__init__(self, user_id=self.id, app_id="", base=base_url)
+    BaseClient.__init__(self, user_id=self.id, app_id="", base=base_url, pat=pat)
     Lister.__init__(self)
 
   def list_apps(self, filter_by: Dict[str, Any] = {}, page_no: int = None,
@@ -57,7 +62,7 @@ class User(Lister, BaseClient):
         per_page=per_page,
         page_no=page_no)
     for app_info in all_apps_info:
-      yield App(base_url=self.base, **app_info)
+      yield App(base_url=self.base, pat=self.pat, **app_info)
 
   def list_runners(self, filter_by: Dict[str, Any] = {}, page_no: int = None,
                    per_page: int = None) -> Generator[Runner, None, None]:
@@ -89,7 +94,7 @@ class User(Lister, BaseClient):
         page_no=page_no)
 
     for runner_info in all_runners_info:
-      yield Runner(check_runner_exists=False, base_url=self.base, **runner_info)
+      yield Runner(check_runner_exists=False, base_url=self.base, pat=self.pat, **runner_info)
 
   def create_app(self, app_id: str, base_workflow: str = 'Empty', **kwargs) -> App:
     """Creates an app for the user.
@@ -115,7 +120,7 @@ class User(Lister, BaseClient):
     if response.status.code != status_code_pb2.SUCCESS:
       raise Exception(response.status)
     self.logger.info("\nApp created\n%s", response.status)
-    kwargs.update({'user_id': self.id, 'base_url': self.base})
+    kwargs.update({'user_id': self.id, 'base_url': self.base, 'pat': self.pat})
     return App(app_id=app_id, **kwargs)
 
   def create_runner(self, runner_id: str, labels: List[str], description: str) -> Runner:
@@ -152,7 +157,8 @@ class User(Lister, BaseClient):
         labels=labels,
         description=description,
         check_runner_exists=False,
-        base_url=self.base)
+        base_url=self.base,
+        pat=self.pat)
 
   def app(self, app_id: str, **kwargs) -> App:
     """Returns an App object for the specified app ID.
@@ -175,7 +181,7 @@ class User(Lister, BaseClient):
       raise Exception(response.status)
 
     kwargs['user_id'] = self.id
-    kwargs.update({'base_url': self.base})
+    kwargs.update({'base_url': self.base, 'pat': self.pat})
     return App(app_id=app_id, **kwargs)
 
   def runner(self, runner_id: str) -> Runner:
@@ -204,7 +210,7 @@ class User(Lister, BaseClient):
     kwargs = self.process_response_keys(dict_response[list(dict_response.keys())[1]],
                                         list(dict_response.keys())[1])
 
-    return Runner(check_runner_exists=False, base_url=self.base, **kwargs)
+    return Runner(check_runner_exists=False, base_url=self.base, pat=self.pat, **kwargs)
 
   def delete_app(self, app_id: str) -> None:
     """Deletes an app for the user.

clarifai/client/workflow.py

@@ -18,16 +18,17 @@ class Workflow(Lister, BaseClient):
   """Workflow is a class that provides access to Clarifai API endpoints related to Workflow information."""
 
   def __init__(self,
-               url_init: str = "",
-               workflow_id: str = "",
+               url: str = None,
+               workflow_id: str = None,
                workflow_version: Dict = {'id': ""},
                output_config: Dict = {'min_value': 0},
                base_url: str = "https://api.clarifai.com",
+               pat: str = None,
                **kwargs):
     """Initializes a Workflow object.
 
     Args:
-        url_init (str): The URL to initialize the workflow object.
+        url (str): The URL to initialize the workflow object.
         workflow_id (str): The Workflow ID to interact with.
         workflow_version (dict): The Workflow Version to interact with.
         output_config (dict): The output config to interact with.
@@ -38,20 +39,20 @@ class Workflow(Lister, BaseClient):
         base_url (str): Base API url. Default "https://api.clarifai.com"
         **kwargs: Additional keyword arguments to be passed to the Workflow.
     """
-    if url_init != "" and workflow_id != "":
-      raise UserError("You can only specify one of url_init or workflow_id.")
-    if url_init == "" and workflow_id == "":
-      raise UserError("You must specify one of url_init or workflow_id.")
-    if url_init != "":
+    if url and workflow_id:
+      raise UserError("You can only specify one of url or workflow_id.")
+    if not url and not workflow_id:
+      raise UserError("You must specify one of url or workflow_id.")
+    if url:
       user_id, app_id, _, workflow_id, workflow_version_id = ClarifaiUrlHelper.split_clarifai_url(
-          url_init)
+          url)
       workflow_version = {'id': workflow_version_id}
       kwargs = {'user_id': user_id, 'app_id': app_id}
     self.kwargs = {**kwargs, 'id': workflow_id, 'version': workflow_version}
     self.output_config = output_config
     self.workflow_info = resources_pb2.Workflow(**self.kwargs)
     self.logger = get_logger(logger_level="INFO")
-    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url)
+    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url, pat=pat)
     Lister.__init__(self)
 
   def predict(self, inputs: List[Input]):
@@ -84,7 +85,7 @@ class Workflow(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.workflow import Workflow
-        >>> workflow = Workflow("workflow_url") # Example: https://clarifai.com/clarifai/main/workflows/Face-Sentiment
+        >>> workflow = Workflow("url") # Example: https://clarifai.com/clarifai/main/workflows/Face-Sentiment
                       or
         >>> workflow = Workflow(user_id='user_id', app_id='app_id', workflow_id='workflow_id')
         >>> workflow_prediction = workflow.predict_by_filepath('filepath', 'image')
@@ -112,13 +113,13 @@ class Workflow(Lister, BaseClient):
       raise UserError('Invalid bytes.')
 
     if input_type == "image":
-      input_proto = Inputs().get_input_from_bytes("", image_bytes=input_bytes)
+      input_proto = Inputs.get_input_from_bytes("", image_bytes=input_bytes)
     elif input_type == "text":
-      input_proto = Inputs().get_input_from_bytes("", text_bytes=input_bytes)
+      input_proto = Inputs.get_input_from_bytes("", text_bytes=input_bytes)
     elif input_type == "video":
-      input_proto = Inputs().get_input_from_bytes("", video_bytes=input_bytes)
+      input_proto = Inputs.get_input_from_bytes("", video_bytes=input_bytes)
     elif input_type == "audio":
-      input_proto = Inputs().get_input_from_bytes("", audio_bytes=input_bytes)
+      input_proto = Inputs.get_input_from_bytes("", audio_bytes=input_bytes)
 
     return self.predict(inputs=[input_proto])
 
@@ -131,7 +132,7 @@ class Workflow(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.workflow import Workflow
-        >>> workflow = Workflow("workflow_url") # Example: https://clarifai.com/clarifai/main/workflows/Face-Sentiment
+        >>> workflow = Workflow("url") # Example: https://clarifai.com/clarifai/main/workflows/Face-Sentiment
                       or
         >>> workflow = Workflow(user_id='user_id', app_id='app_id', workflow_id='workflow_id')
         >>> workflow_prediction = workflow.predict_by_url('url', 'image')
@@ -140,13 +141,13 @@ class Workflow(Lister, BaseClient):
       raise UserError('Invalid input type it should be image, text, video or audio.')
 
     if input_type == "image":
-      input_proto = Inputs().get_input_from_url("", image_url=url)
+      input_proto = Inputs.get_input_from_url("", image_url=url)
     elif input_type == "text":
-      input_proto = Inputs().get_input_from_url("", text_url=url)
+      input_proto = Inputs.get_input_from_url("", text_url=url)
     elif input_type == "video":
-      input_proto = Inputs().get_input_from_url("", video_url=url)
+      input_proto = Inputs.get_input_from_url("", video_url=url)
     elif input_type == "audio":
-      input_proto = Inputs().get_input_from_url("", audio_url=url)
+      input_proto = Inputs.get_input_from_url("", audio_url=url)
 
     return self.predict(inputs=[input_proto])
 
@@ -187,6 +188,7 @@ class Workflow(Lister, BaseClient):
       yield Workflow(
           workflow_id=self.id,
           base_url=self.base,
+          pat=self.pat,
           **dict(self.kwargs, version=workflow_version_info))
 
   def export(self, out_path: str):

clarifai/constants/dataset.py

@@ -0,0 +1,22 @@
+DATASET_UPLOAD_TASKS = [
+    "visual_classification", "text_classification", "visual_detection", "visual_segmentation",
+    "visual_captioning"
+]
+
+TASK_TO_ANNOTATION_TYPE = {
+    "visual_classification": {
+        "concepts": "labels"
+    },
+    "text_classification": {
+        "concepts": "labels"
+    },
+    "visual_captioning": {
+        "concepts": "labels"
+    },
+    "visual_detection": {
+        "bboxes": "bboxes"
+    },
+    "visual_segmentation": {
+        "polygons": "polygons"
+    },
+}

clarifai/datasets/upload/base.py

@@ -3,7 +3,7 @@ from typing import Iterator, List, Tuple, TypeVar, Union
 
 from clarifai_grpc.grpc.api import resources_pb2
 
-from clarifai.client.input import Inputs
+from clarifai.constants.dataset import DATASET_UPLOAD_TASKS
 from clarifai.datasets.upload.features import (TextFeatures, VisualClassificationFeatures,
                                                VisualDetectionFeatures, VisualSegmentationFeatures)
 
@@ -16,18 +16,16 @@ OutputFeaturesType = TypeVar(
 class ClarifaiDataset:
   """Clarifai datasets base class."""
 
-  def __init__(self, datagen_object: Iterator, dataset_id: str, split: str) -> None:
-    self.datagen_object = datagen_object
+  def __init__(self, data_generator: 'ClarifaiDataLoader', dataset_id: str) -> None:
+    self.data_generator = data_generator
     self.dataset_id = dataset_id
-    self.split = split
     self.all_input_ids = {}
     self._all_input_protos = {}
     self._all_annotation_protos = defaultdict(list)
-    self.input_object = Inputs()
 
   def __len__(self) -> int:
     """Get size of all input protos"""
-    return len(self.datagen_object)
+    return len(self.data_generator)
 
   def _to_list(self, input_protos: Iterator) -> List:
     """Parse protos iterator to list."""
@@ -53,9 +51,13 @@ class ClarifaiDataset:
 class ClarifaiDataLoader:
   """Clarifai data loader base class."""
 
-  def __init__(self, split: str) -> None:
+  def __init__(self) -> None:
     pass
 
+  @property
+  def task(self):
+    raise NotImplementedError("Task should be one of {}".format(DATASET_UPLOAD_TASKS))
+
   def load_data(self) -> None:
     raise NotImplementedError()
 

clarifai/datasets/upload/features.py

@@ -16,7 +16,7 @@ class TextFeatures:
 class VisualClassificationFeatures:
   """Image classification datasets preprocessing output features."""
   image_path: str
-  label: Union[str, int]
+  labels: List[Union[str, int]]  # List[str or int] to cater for multi-class tasks
   geo_info: Optional[List[float]] = None  #[Longitude, Latitude]
   id: Optional[int] = None  # image_id
   metadata: Optional[dict] = None
@@ -26,7 +26,7 @@ class VisualClassificationFeatures:
 class VisualDetectionFeatures:
   """Image Detection datasets preprocessing output features."""
   image_path: str
-  classes: List[Union[str, int]]
+  labels: List[Union[str, int]]
   bboxes: List[List[float]]
   geo_info: Optional[List[float]] = None  #[Longitude, Latitude]
   id: Optional[int] = None  # image_id
@@ -37,7 +37,7 @@ class VisualDetectionFeatures:
 class VisualSegmentationFeatures:
   """Image Segmentation datasets preprocessing output features."""
   image_path: str
-  classes: List[Union[str, int]]
+  labels: List[Union[str, int]]
   polygons: List[List[List[float]]]
   geo_info: Optional[List[float]] = None  #[Longitude, Latitude]
   id: Optional[int] = None  # image_id

clarifai/datasets/upload/image.py

@@ -1,17 +1,18 @@
 import os
 from concurrent.futures import ThreadPoolExecutor
-from typing import Iterator, List, Tuple
+from typing import List, Tuple, Type
 
 from clarifai_grpc.grpc.api import resources_pb2
 from google.protobuf.struct_pb2 import Struct
 
-from .base import ClarifaiDataset
+from clarifai.client.input import Inputs
+from clarifai.datasets.upload.base import ClarifaiDataLoader, ClarifaiDataset
 
 
 class VisualClassificationDataset(ClarifaiDataset):
 
-  def __init__(self, datagen_object: Iterator, dataset_id: str, split: str) -> None:
-    super().__init__(datagen_object, dataset_id, split)
+  def __init__(self, data_generator: Type[ClarifaiDataLoader], dataset_id: str) -> None:
+    super().__init__(data_generator, dataset_id)
 
   def _extract_protos(self, batch_input_ids: List[str]
                      ) -> Tuple[List[resources_pb2.Input], List[resources_pb2.Annotation]]:
@@ -24,31 +25,31 @@ class VisualClassificationDataset(ClarifaiDataset):
     """
     input_protos, annotation_protos = [], []
 
-    def process_datagen_item(id):
-      datagen_item = self.datagen_object[id]
+    def process_data_item(id):
+      data_item = self.data_generator[id]
       metadata = Struct()
-      image_path = datagen_item.image_path
-      label = datagen_item.label if isinstance(datagen_item.label,
-                                               list) else [datagen_item.label]  # clarifai concept
-      input_id = f"{self.dataset_id}-{self.split}-{id}" if datagen_item.id is None else f"{self.dataset_id}-{self.split}-{str(datagen_item.id)}"
-      geo_info = datagen_item.geo_info
-      if datagen_item.metadata is not None:
-        metadata.update(datagen_item.metadata)
+      image_path = data_item.image_path
+      labels = data_item.labels if isinstance(data_item.labels,
+                                              list) else [data_item.labels]  # clarifai concept
+      input_id = f"{self.dataset_id}-{id}" if data_item.id is None else f"{self.dataset_id}-{str(data_item.id)}"
+      geo_info = data_item.geo_info
+      if data_item.metadata is not None:
+        metadata.update(data_item.metadata)
       else:
-        metadata.update({"filename": os.path.basename(image_path), "split": self.split})
+        metadata.update({"filename": os.path.basename(image_path)})
 
       self.all_input_ids[id] = input_id
       input_protos.append(
-          self.input_object.get_input_from_file(
+          Inputs.get_input_from_file(
               input_id=input_id,
               image_file=image_path,
               dataset_id=self.dataset_id,
-              labels=label,
+              labels=labels,
               geo_info=geo_info,
               metadata=metadata))
 
     with ThreadPoolExecutor(max_workers=4) as executor:
-      futures = [executor.submit(process_datagen_item, id) for id in batch_input_ids]
+      futures = [executor.submit(process_data_item, id) for id in batch_input_ids]
       for job in futures:
         job.result()
 
@@ -58,8 +59,8 @@ class VisualClassificationDataset(ClarifaiDataset):
 class VisualDetectionDataset(ClarifaiDataset):
   """Visual detection dataset proto class."""
 
-  def __init__(self, datagen_object: Iterator, dataset_id: str, split: str) -> None:
-    super().__init__(datagen_object, dataset_id, split)
+  def __init__(self, data_generator: Type[ClarifaiDataLoader], dataset_id: str) -> None:
+    super().__init__(data_generator, dataset_id)
 
   def _extract_protos(self, batch_input_ids: List[int]
                      ) -> Tuple[List[resources_pb2.Input], List[resources_pb2.Annotation]]:
@@ -72,36 +73,35 @@ class VisualDetectionDataset(ClarifaiDataset):
     """
     input_protos, annotation_protos = [], []
 
-    def process_datagen_item(id):
-      datagen_item = self.datagen_object[id]
+    def process_data_item(id):
+      data_item = self.data_generator[id]
       metadata = Struct()
-      image = datagen_item.image_path
-      labels = datagen_item.classes  # list:[l1,...,ln]
-      bboxes = datagen_item.bboxes  # [[xmin,ymin,xmax,ymax],...,[xmin,ymin,xmax,ymax]]
-      input_id = f"{self.dataset_id}-{self.split}-{id}" if datagen_item.id is None else f"{self.dataset_id}-{self.split}-{str(datagen_item.id)}"
-      if datagen_item.metadata is not None:
-        metadata.update(datagen_item.metadata)
+      image = data_item.image_path
+      labels = data_item.labels  # list:[l1,...,ln]
+      bboxes = data_item.bboxes  # [[xmin,ymin,xmax,ymax],...,[xmin,ymin,xmax,ymax]]
+      input_id = f"{self.dataset_id}-{id}" if data_item.id is None else f"{self.dataset_id}-{str(data_item.id)}"
+      if data_item.metadata is not None:
+        metadata.update(data_item.metadata)
       else:
-        metadata.update({"filename": os.path.basename(image), "split": self.split})
-      geo_info = datagen_item.geo_info
+        metadata.update({"filename": os.path.basename(image)})
+      geo_info = data_item.geo_info
 
       self.all_input_ids[id] = input_id
       input_protos.append(
-          self.input_object.get_input_from_file(
+          Inputs.get_input_from_file(
               input_id=input_id,
               image_file=image,
               dataset_id=self.dataset_id,
               geo_info=geo_info,
               metadata=metadata))
-      # iter over bboxes and classes
+      # iter over bboxes and labels
       # one id could have more than one bbox and label
       for i in range(len(bboxes)):
         annotation_protos.append(
-            self.input_object.get_annotation_proto(
-                input_id=input_id, label=labels[i], annotations=bboxes[i]))
+            Inputs.get_annotation_proto(input_id=input_id, label=labels[i], annotations=bboxes[i]))
 
     with ThreadPoolExecutor(max_workers=4) as executor:
-      futures = [executor.submit(process_datagen_item, id) for id in batch_input_ids]
+      futures = [executor.submit(process_data_item, id) for id in batch_input_ids]
       for job in futures:
         job.result()
 
@@ -111,8 +111,8 @@ class VisualDetectionDataset(ClarifaiDataset):
 class VisualSegmentationDataset(ClarifaiDataset):
   """Visual segmentation dataset proto class."""
 
-  def __init__(self, datagen_object: Iterator, dataset_id: str, split: str) -> None:
-    super().__init__(datagen_object, dataset_id, split)
+  def __init__(self, data_generator: Type[ClarifaiDataLoader], dataset_id: str) -> None:
+    super().__init__(data_generator, dataset_id)
 
   def _extract_protos(self, batch_input_ids: List[str]
                      ) -> Tuple[List[resources_pb2.Input], List[resources_pb2.Annotation]]:
@@ -125,22 +125,22 @@ class VisualSegmentationDataset(ClarifaiDataset):
     """
     input_protos, annotation_protos = [], []
 
-    def process_datagen_item(id):
-      datagen_item = self.datagen_object[id]
+    def process_data_item(id):
+      data_item = self.data_generator[id]
       metadata = Struct()
-      image = datagen_item.image_path
-      labels = datagen_item.classes
-      _polygons = datagen_item.polygons  # list of polygons: [[[x,y],...,[x,y]],...]
-      input_id = f"{self.dataset_id}-{self.split}-{id}" if datagen_item.id is None else f"{self.dataset_id}-{self.split}-{str(datagen_item.id)}"
-      if datagen_item.metadata is not None:
-        metadata.update(datagen_item.metadata)
+      image = data_item.image_path
+      labels = data_item.labels
+      _polygons = data_item.polygons  # list of polygons: [[[x,y],...,[x,y]],...]
+      input_id = f"{self.dataset_id}-{id}" if data_item.id is None else f"{self.dataset_id}-{str(data_item.id)}"
+      if data_item.metadata is not None:
+        metadata.update(data_item.metadata)
       else:
-        metadata.update({"filename": os.path.basename(image), "split": self.split})
-      geo_info = datagen_item.geo_info
+        metadata.update({"filename": os.path.basename(image)})
+      geo_info = data_item.geo_info
 
       self.all_input_ids[id] = input_id
       input_protos.append(
-          self.input_object.get_input_from_file(
+          Inputs.get_input_from_file(
               input_id=input_id,
               image_file=image,
               dataset_id=self.dataset_id,
@@ -152,13 +152,12 @@ class VisualSegmentationDataset(ClarifaiDataset):
       for i, _polygon in enumerate(_polygons):
         try:
           annotation_protos.append(
-              self.input_object.get_mask_proto(
-                  input_id=input_id, label=labels[i], polygons=_polygon))
+              Inputs.get_mask_proto(input_id=input_id, label=labels[i], polygons=_polygon))
         except IndexError:
           continue
 
     with ThreadPoolExecutor(max_workers=4) as executor:
-      futures = [executor.submit(process_datagen_item, id) for id in batch_input_ids]
+      futures = [executor.submit(process_data_item, id) for id in batch_input_ids]
       for job in futures:
         job.result()
 

clarifai/datasets/upload/loaders/coco_captions.py

@@ -1,12 +1,8 @@
-#! COCO 2017 image captioning dataset
+#! COCO image captioning dataset
 
 import os
-import zipfile
-from glob import glob
 
-import requests
 from pycocotools.coco import COCO
-from tqdm import tqdm
 
 from clarifai.datasets.upload.base import ClarifaiDataLoader
 
@@ -14,90 +10,40 @@ from ..features import VisualClassificationFeatures
 
 
 class COCOCaptionsDataLoader(ClarifaiDataLoader):
-  """COCO 2017 Image Captioning Dataset."""
+  """COCO Image Captioning Dataset."""
 
-  def __init__(self, split: str = "train"):
-    """Initialize coco dataset.
+  def __init__(self, images_dir, label_filepath):
+    """
     Args:
-      filenames: the coco zip filenames: Dict[str, str] to be downloaded if download=True,
-      data_dir: the local coco dataset directory.
-      split: "train" or "val"
+      images_dir: Directory containing the images.
+      label_filepath: Path to the COCO annotation file.
     """
-    self.filenames = {
-        "train": "train2017.zip",
-        "val": "val2017.zip",
-        "annotations": "annotations_trainval2017.zip"
-    }
-    self.split = split
-    self.url = "http://images.cocodataset.org/zips/"  # coco base image-zip url
-    self.data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
-                                 "data")  # data storage directory
-    self.extracted_coco_dirs = {"train": None, "val": None, "annotations": None}
+    self.images_dir = images_dir
+    self.label_filepath = label_filepath
 
+    self.map_ids = {}
     self.load_data()
 
-  def coco_download(self, save_dir):
-    """Download coco dataset."""
-    if not os.path.exists(save_dir):
-      os.mkdir(save_dir)
-
-    #check if train, val and annotation dirs exist
-    #so that the coco2017 data isn't downloaded
-    for key, filename in self.filenames.items():
-      existing_files = glob(f"{save_dir}/{key}*")
-      if existing_files:
-        print(f"{key} dataset already downloded and extracted")
-        continue
-
-      print("-" * 80)
-      print(f"Downloading {filename}")
-      print("-" * 80)
-
-      if "annotations" in filename:
-        self.url = "http://images.cocodataset.org/annotations/"
+  @property
+  def task(self):
+    return "visual_captioning"
 
-      response = requests.get(self.url + filename, stream=True)
-      response.raise_for_status()
-      with open(os.path.join(save_dir, filename), "wb") as _file:
-        for chunk in tqdm(response.iter_content(chunk_size=5124000)):
-          if chunk:
-            _file.write(chunk)
-      print("Data download complete...")
+  def load_data(self) -> None:
+    self.coco = COCO(self.label_filepath)
+    self.map_ids = {i: img_id for i, img_id in enumerate(list(self.coco.imgs.keys()))}
 
-      #extract files
-      zf = zipfile.ZipFile(os.path.join(save_dir, filename))
-      print(f" Extracting {filename} file")
-      zf.extractall(path=save_dir)
-      # Delete coco zip
-      print(f" Deleting {filename}")
-      os.remove(path=os.path.join(save_dir, filename))
-
-  def load_data(self):
-    if isinstance(self.filenames, dict) and len(self.filenames) == 3:
-      self.coco_download(self.data_dir)
-      self.extracted_coco_dirs["train"] = [os.path.join(self.data_dir, i) \
-      for i in os.listdir(self.data_dir) if "train" in i][0]
-      self.extracted_coco_dirs["val"] = [os.path.join(self.data_dir, i) \
-      for i in os.listdir(self.data_dir) if "val" in i][0]
-
-      self.extracted_coco_dirs["annotations"] = [os.path.join(self.data_dir, i) \
-      for i in os.listdir(self.data_dir) if "annotations" in i][0]
-    else:
-      raise Exception(f"`filenames` must be a dict of atleast 2 coco zip file names; \
-  train, val and annotations. Found {len(self.filenames)} items instead.")
+  def __len__(self):
+    return len(self.coco.imgs)
 
-    annot_file = glob(self.extracted_coco_dirs["annotations"] + "/" + f"captions_{self.split}*")[0]
-    coco = COCO(annot_file)
-    annot_ids = coco.getAnnIds()
-    self.annotations = coco.loadAnns(annot_ids)
+  def __getitem__(self, index):
+    value = self.coco.imgs[self.map_ids[index]]
+    image_path = os.path.join(self.images_dir, value['file_name'])
+    annots = []
 
-  def __len__(self):
-    return len(self.annotations)
+    input_ann_ids = self.coco.getAnnIds(imgIds=[value['id']])
+    input_anns = self.coco.loadAnns(input_ann_ids)
 
-  def __getitem__(self, idx):
-    annot = self.annotations[idx]
-    image_path = glob(
-        os.path.join(self.extracted_coco_dirs[self.split],
-                     f"{str(annot['image_id']).zfill(12)}*"))[0]
+    for ann in input_anns:
+      annots.append(ann['caption'])
 
-    return VisualClassificationFeatures(image_path, annot["caption"], id=annot["image_id"])
+    return VisualClassificationFeatures(image_path, labels=annots[0], id=str(value['id']))