No public description

PiperOrigin-RevId: 569017854

Author: tensorflower-gardener

official/projects/waste_identification_ml/two_model_inference/download_and_unzip_models.py

@@ -0,0 +1,100 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""This module provides utilities for executing shell commands.
+
+It particularly downloads and extracts Mask RCNN models from the TensorFlow
+model garden. It includes a function to execute shell commands and
+a custom exception to handle errors that arise from command execution.
+
+Functions:
+    - execute_command(cmd: str) -> str: Executes a shell command and returns its
+    standard output. Raises
+      a CommandExecutionError if the command execution fails.
+
+Exceptions:
+    - CommandExecutionError: Custom exception that's raised when there's an
+    error executing a shell command.
+
+Usage:
+    The main purpose of this module is to download two specific Mask RCNN models
+    and unzip them. The module
+    performs these operations when imported.
+
+Note:
+    It's recommended to not perform actions like downloading files on module
+    import in production applications.
+    It's better to move such tasks inside a function or a main block to allow
+    for more controlled execution.
+"""
+import argparse
+import os
+import subprocess
+
+
+class CommandExecutionError(Exception):
+  """Raised when there's an error executing a shell command."""
+
+  def __init__(self, cmd, returncode, stderr):
+    super().__init__(f"Error executing command: {cmd}. Error: {stderr}")
+    self.cmd = cmd
+    self.returncode = returncode
+    self.stderr = stderr
+
+
+def execute_command(cmd: str) -> str:
+  """Executes a shell command and returns its output."""
+  result = subprocess.run(
+      cmd,
+      shell=True,
+      stdout=subprocess.PIPE,
+      stderr=subprocess.PIPE,
+      check=False,
+  )
+
+  if result.returncode != 0:
+    raise CommandExecutionError(
+        cmd, result.returncode, result.stderr.decode("utf-8")
+    )
+
+  return result.stdout.decode("utf-8")
+
+
+def main(_) -> None:
+  # Download the provided files
+  execute_command(f"wget {args.url1}")
+  execute_command(f"wget {args.url2}")
+
+  # Create directories
+  os.makedirs("material", exist_ok=True)
+  os.makedirs("material_form", exist_ok=True)
+
+  # Unzip the provided files
+  zip_file1 = os.path.basename(args.url1)
+  zip_file2 = os.path.basename(args.url2)
+  execute_command(f"unzip {zip_file1} -d material/")
+  execute_command(f"unzip {zip_file2} -d material_form/")
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(
+      description="Download and extract Mask RCNN models."
+  )
+  parser.add_argument("material_url", help="repo url for material model")
+  parser.add_argument(
+      "material_form_url", help="repo url for material form model"
+  )
+
+  args = parser.parse_args()
+  main(args)

official/projects/waste_identification_ml/two_model_inference/labels.py

@@ -0,0 +1,118 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Load labels for model prediction.
+
+Given paths of CSV files, task is to import them and convert into a
+form required for mapping with the model output.
+"""
+import csv
+from typing import TypedDict
+
+
+class ItemDict(TypedDict):
+  id: int
+  name: str
+  supercategory: str
+
+
+def read_csv_to_list(file_path: str) -> list[str]:
+  """Reads a CSV file and returns its contents as a list.
+
+  This function reads the given CSV file, skips the header, and assumes
+  there is only one column in the CSV. It returns the contents as a list of
+  strings.
+
+  Args:
+      file_path: The path to the CSV file.
+
+  Returns:
+      The contents of the CSV file as a list of strings.
+  """
+  data_list = []
+  with open(file_path, 'r') as csvfile:
+    reader = csv.reader(csvfile)
+    next(reader)  # Skip the header row if present
+    for row in reader:
+      data_list.append(row[0])  # Assuming there is only one column in the CSV
+  return data_list
+
+
+def categories_dictionary(objects: list[str]) -> dict[int, ItemDict]:
+  """This function takes a list of objects and returns a dictionaries.
+
+  A dictionary of objects, where each object is represented by a dictionary
+  with the following keys:
+    - id: The ID of the object.
+    - name: The name of the object.
+    - supercategory: The supercategory of the object.
+
+  Args:
+    objects: A list of strings, where each string is the name of an
+      object.
+
+  Returns:
+    A tuple of two dictionaries, as described above.
+  """
+  category_index = {}
+
+  for num, obj_name in enumerate(objects, start=1):
+    obj_dict = {'id': num, 'name': obj_name, 'supercategory': 'objects'}
+    category_index[num] = obj_dict
+
+  return category_index
+
+
+def load_labels(
+    label_paths: dict[str, str]
+) -> tuple[list[list[str]], dict[int, ItemDict]]:
+  """Loads labels, combines them, and formats them for prediction.
+
+  This function reads labels for multiple models, combines the labels in
+  order to predict a single label output, and formats them into the desired
+  structure required for prediction.
+
+  Args:
+    label_paths: Dictionary of label paths for different models.
+
+  Returns:
+          - A list of lists containing individual category indices for each
+          model.
+          - A dictionary of combined category indices in the desired format for
+          prediction.
+
+  Note:
+      - The function assumes there are exactly two models.
+      - Inserts a category 'Na' for both models in case there is no detection.
+      - The total number of predicted labels for a combined model is
+      predetermined.
+  """
+  # loading labels for both models
+  category_indices = [read_csv_to_list(label) for label in label_paths.values()]
+
+  # insert a cateory 'Na' for both models in case there is no detection
+  for i in [0, 1]:
+    category_indices[i].insert(0, 'Na')
+
+  # combine the labels for both models in order to predict a single label output
+  combined_category_indices = []
+  for i in category_indices[0]:
+    for j in category_indices[1]:
+      combined_category_indices.append(f'{i}_{j}')
+  combined_category_indices.sort()
+
+  # convert the list of labels into a desired format required for prediction
+  category_index = categories_dictionary(combined_category_indices)
+
+  return category_indices, category_index