megadetector 5.0.11__py3-none-any.whl → 5.0.12__py3-none-any.whl

megadetector/api/batch_processing/api_core/batch_service/score.py

@@ -0,0 +1,439 @@
+import io
+import json
+import math
+import os
+import sys
+from datetime import datetime
+from io import BytesIO
+from typing import Union
+
+from PIL import Image
+import numpy as np
+import requests
+import tensorflow as tf
+from azure.storage.blob import ContainerClient
+
+print('score.py, tensorflow version:', tf.__version__)
+print('score.py, tf.test.is_gpu_available:', tf.test.is_gpu_available())
+
+PRINT_EVERY = 500
+
+
+#%% Helper functions *copied* from ct_utils.py and visualization/visualization_utils.py
+
+IMAGE_ROTATIONS = {
+    3: 180,
+    6: 270,
+    8: 90
+}
+
+def truncate_float(x, precision=3):
+    """
+    Function for truncating a float scalar to the defined precision.
+    For example: truncate_float(0.0003214884) --> 0.000321
+    This function is primarily used to achieve a certain float representation
+    before exporting to JSON
+    Args:
+    x         (float) Scalar to truncate
+    precision (int)   The number of significant digits to preserve, should be
+                      greater or equal 1
+    """
+
+    assert precision > 0
+
+    if np.isclose(x, 0):
+        return 0
+    else:
+        # Determine the factor, which shifts the decimal point of x
+        # just behind the last significant digit
+        factor = math.pow(10, precision - 1 - math.floor(math.log10(abs(x))))
+        # Shift decimal point by multiplicatipon with factor, flooring, and
+        # division by factor
+        return math.floor(x * factor)/factor
+
+
+def open_image(input_file: Union[str, BytesIO]) -> Image:
+    """Opens an image in binary format using PIL.Image and converts to RGB mode.
+
+    This operation is lazy; image will not be actually loaded until the first
+    operation that needs to load it (for example, resizing), so file opening
+    errors can show up later.
+
+    Args:
+        input_file: str or BytesIO, either a path to an image file (anything
+            that PIL can open), or an image as a stream of bytes
+
+    Returns:
+        an PIL image object in RGB mode
+    """
+    if (isinstance(input_file, str)
+            and input_file.startswith(('http://', 'https://'))):
+        response = requests.get(input_file)
+        image = Image.open(BytesIO(response.content))
+        try:
+            response = requests.get(input_file)
+            image = Image.open(BytesIO(response.content))
+        except Exception as e:
+            print(f'Error opening image {input_file}: {e}')
+            raise
+    else:
+        image = Image.open(input_file)
+    if image.mode not in ('RGBA', 'RGB', 'L'):
+        raise AttributeError(f'Image {input_file} uses unsupported mode {image.mode}')
+    if image.mode == 'RGBA' or image.mode == 'L':
+        # PIL.Image.convert() returns a converted copy of this image
+        image = image.convert(mode='RGB')
+
+    # alter orientation as needed according to EXIF tag 0x112 (274) for Orientation
+    # https://gist.github.com/dangtrinhnt/a577ece4cbe5364aad28
+    # https://www.media.mit.edu/pia/Research/deepview/exif.html
+    try:
+        exif = image._getexif()
+        orientation: int = exif.get(274, None)  # 274 is the key for the Orientation field
+        if orientation is not None and orientation in IMAGE_ROTATIONS:
+            image = image.rotate(IMAGE_ROTATIONS[orientation], expand=True)  # returns a rotated copy
+    except Exception:
+        pass
+
+    return image
+
+
+def load_image(input_file: Union[str, BytesIO]) -> Image.Image:
+    """Loads the image at input_file as a PIL Image into memory.
+    Image.open() used in open_image() is lazy and errors will occur downstream
+    if not explicitly loaded.
+    Args:
+        input_file: str or BytesIO, either a path to an image file (anything
+            that PIL can open), or an image as a stream of bytes
+    Returns: PIL.Image.Image, in RGB mode
+    """
+    image = open_image(input_file)
+    image.load()
+    return image
+
+
+#%% TFDetector class, an unmodified *copy* of the class in detection/tf_detector.py,
+# so we do not have to import the packages required by run_detector.py
+
+class TFDetector:
+    """
+    A detector model loaded at the time of initialization. It is intended to be used with
+    MegaDetector (TF). The inference batch size is set to 1; code needs to be modified
+    to support larger batch sizes, including resizing appropriately.
+    """
+
+    # Number of decimal places to round to for confidence and bbox coordinates
+    CONF_DIGITS = 3
+    COORD_DIGITS = 4
+
+    # MegaDetector was trained with batch size of 1, and the resizing function is a part
+    # of the inference graph
+    BATCH_SIZE = 1
+
+    # An enumeration of failure reasons
+    FAILURE_TF_INFER = 'Failure TF inference'
+    FAILURE_IMAGE_OPEN = 'Failure image access'
+
+    DEFAULT_RENDERING_CONFIDENCE_THRESHOLD = 0.85  # to render bounding boxes
+    DEFAULT_OUTPUT_CONFIDENCE_THRESHOLD = 0.1  # to include in the output json file
+
+    DEFAULT_DETECTOR_LABEL_MAP = {
+        '1': 'animal',
+        '2': 'person',
+        '3': 'vehicle'  # available in megadetector v4+
+    }
+
+    NUM_DETECTOR_CATEGORIES = 4  # animal, person, group, vehicle - for color assignment
+
+    def __init__(self, model_path):
+        """Loads model from model_path and starts a tf.Session with this graph. Obtains
+        input and output tensor handles."""
+        detection_graph = TFDetector.__load_model(model_path)
+        self.tf_session = tf.Session(graph=detection_graph)
+
+        self.image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
+        self.box_tensor = detection_graph.get_tensor_by_name('detection_boxes:0')
+        self.score_tensor = detection_graph.get_tensor_by_name('detection_scores:0')
+        self.class_tensor = detection_graph.get_tensor_by_name('detection_classes:0')
+
+    @staticmethod
+    def round_and_make_float(d, precision=4):
+        return truncate_float(float(d), precision=precision)
+
+    @staticmethod
+    def __convert_coords(tf_coords):
+        """Converts coordinates from the model's output format [y1, x1, y2, x2] to the
+        format used by our API and MegaDB: [x1, y1, width, height]. All coordinates
+        (including model outputs) are normalized in the range [0, 1].
+        Args:
+            tf_coords: np.array of predicted bounding box coordinates from the TF detector,
+                has format [y1, x1, y2, x2]
+        Returns: list of Python float, predicted bounding box coordinates [x1, y1, width, height]
+        """
+        # change from [y1, x1, y2, x2] to [x1, y1, width, height]
+        width = tf_coords[3] - tf_coords[1]
+        height = tf_coords[2] - tf_coords[0]
+
+        new = [tf_coords[1], tf_coords[0], width, height]  # must be a list instead of np.array
+
+        # convert numpy floats to Python floats
+        for i, d in enumerate(new):
+            new[i] = TFDetector.round_and_make_float(d, precision=TFDetector.COORD_DIGITS)
+        return new
+
+    @staticmethod
+    def convert_to_tf_coords(array):
+        """From [x1, y1, width, height] to [y1, x1, y2, x2], where x1 is x_min, x2 is x_max
+        This is an extraneous step as the model outputs [y1, x1, y2, x2] but were converted to the API
+        output format - only to keep the interface of the sync API.
+        """
+        x1 = array[0]
+        y1 = array[1]
+        width = array[2]
+        height = array[3]
+        x2 = x1 + width
+        y2 = y1 + height
+        return [y1, x1, y2, x2]
+
+    @staticmethod
+    def __load_model(model_path):
+        """Loads a detection model (i.e., create a graph) from a .pb file.
+        Args:
+            model_path: .pb file of the model.
+        Returns: the loaded graph.
+        """
+        print('TFDetector: Loading graph...')
+        detection_graph = tf.Graph()
+        with detection_graph.as_default():
+            od_graph_def = tf.GraphDef()
+            with tf.gfile.GFile(model_path, 'rb') as fid:
+                serialized_graph = fid.read()
+                od_graph_def.ParseFromString(serialized_graph)
+                tf.import_graph_def(od_graph_def, name='')
+        print('TFDetector: Detection graph loaded.')
+
+        return detection_graph
+
+    def _generate_detections_one_image(self, image):
+        np_im = np.asarray(image, np.uint8)
+        im_w_batch_dim = np.expand_dims(np_im, axis=0)
+
+        # need to change the above line to the following if supporting a batch size > 1 and resizing to the same size
+        # np_images = [np.asarray(image, np.uint8) for image in images]
+        # images_stacked = np.stack(np_images, axis=0) if len(images) > 1 else np.expand_dims(np_images[0], axis=0)
+
+        # performs inference
+        (box_tensor_out, score_tensor_out, class_tensor_out) = self.tf_session.run(
+            [self.box_tensor, self.score_tensor, self.class_tensor],
+            feed_dict={self.image_tensor: im_w_batch_dim})
+
+        return box_tensor_out, score_tensor_out, class_tensor_out
+
+    def generate_detections_one_image(self, image, image_id,
+                                      detection_threshold=DEFAULT_OUTPUT_CONFIDENCE_THRESHOLD):
+        """Apply the detector to an image.
+        Args:
+            image: the PIL Image object
+            image_id: a path to identify the image; will be in the "file" field of the output object
+            detection_threshold: confidence above which to include the detection proposal
+        Returns:
+        A dict with the following fields, see the 'images' key in https://github.com/agentmorris/MegaDetector/tree/main/megadetector/api/batch_processing#batch-processing-api-output-format
+            - 'file' (always present)
+            - 'max_detection_conf'
+            - 'detections', which is a list of detection objects containing keys 'category', 'conf' and 'bbox'
+            - 'failure'
+        """
+        result = {
+            'file': image_id
+        }
+        try:
+            b_box, b_score, b_class = self._generate_detections_one_image(image)
+
+            # our batch size is 1; need to loop the batch dim if supporting batch size > 1
+            boxes, scores, classes = b_box[0], b_score[0], b_class[0]
+
+            detections_cur_image = []  # will be empty for an image with no confident detections
+            max_detection_conf = 0.0
+            for b, s, c in zip(boxes, scores, classes):
+                if s > detection_threshold:
+                    detection_entry = {
+                        'category': str(int(c)),  # use string type for the numerical class label, not int
+                        'conf': truncate_float(float(s),  # cast to float for json serialization
+                                               precision=TFDetector.CONF_DIGITS),
+                        'bbox': TFDetector.__convert_coords(b)
+                    }
+                    detections_cur_image.append(detection_entry)
+                    if s > max_detection_conf:
+                        max_detection_conf = s
+
+            result['max_detection_conf'] = truncate_float(float(max_detection_conf),
+                                                          precision=TFDetector.CONF_DIGITS)
+            result['detections'] = detections_cur_image
+
+        except Exception as e:
+            result['failure'] = TFDetector.FAILURE_TF_INFER
+            print('TFDetector: image {} failed during inference: {}'.format(image_id, str(e)))
+
+        return result
+
+
+#%% Scoring script
+
+class BatchScorer:
+    """
+    Coordinates scoring images in this Task.
+
+    1. have a synchronized queue that download tasks enqueue and scoring function dequeues - but need to be able to
+    limit the size of the queue. We do not want to write the image to disk and then load it in the scoring func.
+    """
+    def __init__(self, **kwargs):
+        print('score.py BatchScorer, __init__()')
+
+        detector_path = kwargs.get('detector_path')
+        self.detector = TFDetector(detector_path)
+
+        self.use_url = kwargs.get('use_url')
+        if not self.use_url:
+            input_container_sas = kwargs.get('input_container_sas')
+            self.input_container_client = ContainerClient.from_container_url(input_container_sas)
+
+        self.detection_threshold = kwargs.get('detection_threshold')
+
+        self.image_ids_to_score = kwargs.get('image_ids_to_score')
+
+        # determine if there is metadata attached to each image_id
+        self.metadata_available = True if isinstance(self.image_ids_to_score[0], list) else False
+
+    def _download_image(self, image_file) -> Image:
+        """
+        Args:
+            image_file: Public URL if use_url, else the full path from container root
+
+        Returns:
+            PIL image loaded
+        """
+        if not self.use_url:
+            downloader = self.input_container_client.download_blob(image_file)
+            image_file = io.BytesIO()
+            blob_props = downloader.download_to_stream(image_file)
+
+        image = open_image(image_file)
+        return image
+
+    def score_images(self) -> list:
+        detections = []
+
+        for i in self.image_ids_to_score:
+
+            if self.metadata_available:
+                image_id = i[0]
+                image_metadata = i[1]
+            else:
+                image_id = i
+
+            try:
+                image = self._download_image(image_id)
+            except Exception as e:
+                print(f'score.py BatchScorer, score_images, download_image exception: {e}')
+                result = {
+                    'file': image_id,
+                    'failure': TFDetector.FAILURE_IMAGE_OPEN
+                }
+            else:
+                result = self.detector.generate_detections_one_image(
+                    image, image_id, detection_threshold=self.detection_threshold)
+
+            if self.metadata_available:
+                result['meta'] = image_metadata
+
+            detections.append(result)
+            if len(detections) % PRINT_EVERY == 0:
+                print(f'scored {len(detections)} images')
+
+        return detections
+
+
+def main():
+    print('score.py, main()')
+
+    # information to determine input and output locations
+    api_instance_name = os.environ['API_INSTANCE_NAME']
+    job_id = os.environ['AZ_BATCH_JOB_ID']
+    task_id = os.environ['AZ_BATCH_TASK_ID']
+    mount_point = os.environ['AZ_BATCH_NODE_MOUNTS_DIR']
+
+    # other parameters for the task
+    begin_index = int(os.environ['TASK_BEGIN_INDEX'])
+    end_index = int(os.environ['TASK_END_INDEX'])
+
+    input_container_sas = os.environ.get('JOB_CONTAINER_SAS', None)  # could be None if use_url
+    use_url = os.environ.get('JOB_USE_URL', None)
+
+    if use_url and use_url.lower() == 'true':  # bool of any non-empty string is True
+        use_url = True
+    else:
+        use_url = False
+
+    detection_threshold = float(os.environ['DETECTION_CONF_THRESHOLD'])
+
+    print(f'score.py, main(), api_instance_name: {api_instance_name}, job_id: {job_id}, task_id: {task_id}, '
+          f'mount_point: {mount_point}, begin_index: {begin_index}, end_index: {end_index}, '
+          f'input_container_sas: {input_container_sas}, use_url (parsed): {use_url}'
+          f'detection_threshold: {detection_threshold}')
+
+    job_folder_mounted = os.path.join(mount_point, 'batch-api', f'api_{api_instance_name}', f'job_{job_id}')
+    task_out_dir = os.path.join(job_folder_mounted, 'task_outputs')
+    os.makedirs(task_out_dir, exist_ok=True)
+    task_output_path = os.path.join(task_out_dir, f'job_{job_id}_task_{task_id}.json')
+
+    # test that we can write to output path; also in case there is no image to process
+    with open(task_output_path, 'w') as f:
+        json.dump([], f)
+
+    # list images to process
+    list_images_path = os.path.join(job_folder_mounted, f'{job_id}_images.json')
+    with open(list_images_path) as f:
+        list_images = json.load(f)
+    print(f'score.py, main(), length of list_images: {len(list_images)}')
+
+    if (not isinstance(list_images, list)) or len(list_images) == 0:
+        print('score.py, main(), zero images in specified overall list, exiting...')
+        sys.exit(0)
+
+    # items in this list can be strings or [image_id, metadata]
+    list_images = list_images[begin_index: end_index]
+    if len(list_images) == 0:
+        print('score.py, main(), zero images in the shard, exiting')
+        sys.exit(0)
+
+    print(f'score.py, main(), processing {len(list_images)} images in this Task')
+
+    # model path
+    # Path to .pb TensorFlow detector model file, relative to the
+    # models/megadetector_copies folder in mounted container
+    detector_model_rel_path = os.environ['DETECTOR_REL_PATH']
+    detector_path = os.path.join(mount_point, 'models', 'megadetector_copies', detector_model_rel_path)
+    assert os.path.exists(detector_path), f'detector is not found at the specified path: {detector_path}'
+
+    # score the images
+    scorer = BatchScorer(
+        detector_path=detector_path,
+        use_url=use_url,
+        input_container_sas=input_container_sas,
+        detection_threshold=detection_threshold,
+        image_ids_to_score=list_images
+    )
+
+    try:
+        tick = datetime.now()
+        detections = scorer.score_images()
+        duration = datetime.now() - tick
+        print(f'score.py, main(), score_images() duration: {duration}')
+    except Exception as e:
+        raise RuntimeError(f'score.py, main(), exception in score_images(): {e}')
+
+    with open(task_output_path, 'w', encoding='utf-8') as f:
+        json.dump(detections, f, ensure_ascii=False)
+
+if __name__ == '__main__':
+    main()