Post process

Post-processing functions.

detection_folder_separation(json_file, img_path, destination_path, confidence_threshold)

Processes detection data from a JSON file to sort images into 'Animal' or 'No_animal' directories based on detection categories and confidence levels.

This function reads a JSON formatted file containing annotations of image detections. Each image is checked for detections with category '0' and a confidence level above the specified threshold. If such detections are found, the image is categorized under 'Animal'. Images without any category '0' detections above the threshold, including those with no detections at all, are categorized under 'No_animal'.

Parameters: - json_file (str): Path to the JSON file containing detection data. - destination_path (str): Base path where 'Animal' and 'No_animal' folders will be created and into which images will be sorted and copied. - source_images_directory (str): Path to the directory containing the source images to be processed. - confidence_threshold (float): The confidence threshold to consider a detection as valid.

Effects: - Reads from the specified json_file. - Copies files from source_images_directory to either destination_path/Animal or destination_path/No_animal based on the detection data and confidence level.

Note: - The function assumes that the JSON file structure includes keys 'annotations', each containing 'img_id', 'bbox', 'category', and 'confidence'. It does not handle missing keys or unexpected JSON structures and may raise an exception in such cases. - Directories Animal and No_animal are created if they do not already exist. - Images are copied, not moved; original images remain in the source directory.

Source code in PytorchWildlife/utils/post_process.py

def detection_folder_separation(json_file, img_path, destination_path, confidence_threshold):
    """
    Processes detection data from a JSON file to sort images into 'Animal' or 'No_animal' directories
    based on detection categories and confidence levels.

    This function reads a JSON formatted file containing annotations of image detections.
    Each image is checked for detections with category '0' and a confidence level above the specified
    threshold. If such detections are found, the image is categorized under 'Animal'. Images without
    any category '0' detections above the threshold, including those with no detections at all, are 
    categorized under 'No_animal'.

    Parameters:
    - json_file (str): Path to the JSON file containing detection data.
    - destination_path (str): Base path where 'Animal' and 'No_animal' folders will be created
                              and into which images will be sorted and copied.
    - source_images_directory (str): Path to the directory containing the source images to be processed.
    - confidence_threshold (float): The confidence threshold to consider a detection as valid.

    Effects:
    - Reads from the specified `json_file`.
    - Copies files from `source_images_directory` to either `destination_path/Animal` or
      `destination_path/No_animal` based on the detection data and confidence level.

    Note:
    - The function assumes that the JSON file structure includes keys 'annotations', each containing
      'img_id', 'bbox', 'category', and 'confidence'. It does not handle missing keys or unexpected
      JSON structures and may raise an exception in such cases.
    - Directories `Animal` and `No_animal` are created if they do not already exist.
    - Images are copied, not moved; original images remain in the source directory.
    """

    # Load JSON data from the file
    with open(json_file, 'r') as file:
        data = json.load(file)

    # Ensure the destination directories exist
    os.makedirs(destination_path, exist_ok=True)
    animal_path = os.path.join(destination_path, "Animal")
    no_animal_path = os.path.join(destination_path, "No_animal")
    os.makedirs(animal_path, exist_ok=True)
    os.makedirs(no_animal_path, exist_ok=True)

    # Process each image detection
    i = 0
    for item in data['annotations']:
        i+=1
        img_id = item['img_id']
        categories = item['category']
        confidences = item['confidence']

        # Check if there is any category '0' with confidence above the threshold
        file_targeted_for_animal = False
        for category, confidence in zip(categories, confidences):
            if category == 0 and confidence > confidence_threshold:
                file_targeted_for_animal = True
                break

        if file_targeted_for_animal:
            target_folder = animal_path
        else:
            target_folder = no_animal_path

        # Construct the source and destination file paths
        src_file_path = os.path.join(img_path, img_id)
        dest_file_path = os.path.join(target_folder, os.path.dirname(img_id))
        os.makedirs(dest_file_path, exist_ok=True)

        # Copy the file to the appropriate directory
        shutil.copy(src_file_path, dest_file_path)

    return "{} files were successfully separated".format(i)

save_crop_images(results, output_dir, input_dir=None, overwrite=False)

Save cropped images based on the detection bounding boxes.

Parameters:

Name	Type	Description	Default
results	list	Detection results containing image ID and detections.	required
output_dir	str	Directory to save the cropped images.	required
input_dir	str	Directory containing the input images. Default to None.	None
overwrite	bool	Whether overwriting existing image folders. Default to False.	False

Source code in PytorchWildlife/utils/post_process.py

def save_crop_images(results, output_dir, input_dir=None, overwrite=False):
    """
    Save cropped images based on the detection bounding boxes.

    Args:
        results (list):
            Detection results containing image ID and detections.
        output_dir (str):
            Directory to save the cropped images.
        input_dir (str):
            Directory containing the input images. Default to None.
        overwrite (bool):
            Whether overwriting existing image folders. Default to False.
    """

    os.makedirs(output_dir, exist_ok=True)

    with sv.ImageSink(target_dir_path=output_dir, overwrite=overwrite) as sink:
        if isinstance(results, list):
            for entry in results:
                for i, (xyxy, cat) in enumerate(zip(entry["detections"].xyxy, entry["detections"].class_id)):
                    cropped_img = sv.crop_image(
                        image=np.array(Image.open(entry["img_id"]).convert("RGB")), xyxy=xyxy
                    )
                    if input_dir:
                        relative_path = os.path.relpath(entry["img_id"], input_dir)
                        save_path = os.path.join(output_dir, relative_path)
                        os.makedirs(os.path.dirname(save_path), exist_ok=True) 
                        image_name = os.path.join(os.path.dirname(relative_path), "{}_{}_{}".format(int(cat), i, os.path.basename(entry["img_id"])))
                    else:
                        image_name = "{}_{}_{}".format(int(cat), i, os.path.basename(entry["img_id"]))
                    sink.save_image(
                        image=cv2.cvtColor(cropped_img, cv2.COLOR_RGB2BGR),
                        image_name=image_name,
                    )
        else:
            for i, (xyxy, cat) in enumerate(zip(results["detections"].xyxy, results["detections"].class_id)):
                cropped_img = sv.crop_image(
                    image=np.array(Image.open(results["img_id"]).convert("RGB")), xyxy=xyxy
                )
                sink.save_image(
                    image=cv2.cvtColor(cropped_img, cv2.COLOR_RGB2BGR),
                    image_name="{}_{}_{}".format(int(cat), i, os.path.basename(results["img_id"]),
                ))

save_detection_classification_json(det_results, clf_results, output_path, det_categories=None, clf_categories=None, exclude_file_path=None)

Save classification results to a JSON file.

Parameters:

Name	Type	Description	Default
det_results	list	Detection results containing image ID, bounding boxes, detection category, and confidence.	required
clf_results	list	classification results containing image ID, classification category, and confidence.	required
output_path	str	Path to save the output JSON file.	required
det_categories	list	List of categories for detected objects. Defaults to None.	None
clf_categories	list	List of categories for classified objects. Defaults to None.	None
exclude_file_path	str	We can exclude the some path sections from the image ID. Defaults to None.	None

Source code in PytorchWildlife/utils/post_process.py

def save_detection_classification_json(
    det_results, clf_results, output_path, det_categories=None, clf_categories=None, exclude_file_path=None
):
    """
    Save classification results to a JSON file.

    Args:
        det_results (list):
            Detection results containing image ID, bounding boxes, detection category, and confidence.
        clf_results (list):
            classification results containing image ID, classification category, and confidence.
        output_path (str):
            Path to save the output JSON file.
        det_categories (list, optional):
            List of categories for detected objects. Defaults to None.
        clf_categories (list, optional):
            List of categories for classified objects. Defaults to None.
        exclude_file_path (str, optional):
            We can exclude the some path sections from the image ID. Defaults to None.
    """

    json_results = {
        "annotations": [],
        "det_categories": det_categories,
        "clf_categories": clf_categories,
    }

    with open(output_path, "w") as f:
        counter = 0
        for det_r in det_results:
            clf_categories = []
            clf_confidence = []
            for i in range(counter, len(clf_results)):
                clf_r = clf_results[i]
                if clf_r["img_id"] == det_r["img_id"]:
                    clf_categories.append(clf_r["class_id"])
                    clf_confidence.append(clf_r["confidence"])
                    counter += 1
                else:
                    break

            json_results["annotations"].append(
                {
                    "img_id": str(det_r["img_id"]).replace(exclude_file_path + os.sep, '') if exclude_file_path else str(det_r["img_id"]),
                    "bbox": [
                        [int(x) for x in sublist]
                        for sublist in det_r["detections"].xyxy.astype(int).tolist()
                    ],
                    "det_category": [
                        int(x) for x in det_r["detections"].class_id.tolist()
                    ],
                    "det_confidence": [
                        float(x) for x in det_r["detections"].confidence.tolist()
                    ],
                    "clf_category": [int(x) for x in clf_categories],
                    "clf_confidence": [float(x) for x in clf_confidence],
                }
            )
        json.dump(json_results, f, indent=4)

save_detection_classification_timelapse_json(det_results, clf_results, output_path, det_categories=None, clf_categories=None, exclude_file_path=None, info={'detector': 'megadetector_v5'})

Save detection and classification results to a JSON file in the specified format.

Parameters:

Name	Type	Description	Default
det_results	list	Detection results containing image ID, bounding boxes, detection category, and confidence.	required
clf_results	list	Classification results containing image ID, classification category, and confidence.	required
output_path	str	Path to save the output JSON file.	required
det_categories	dict	Dictionary of categories for detected objects. Defaults to None.	None
clf_categories	dict	Dictionary of categories for classified objects. Defaults to None.	None
exclude_file_path	str	We can exclude the some path sections from the image ID. Defaults to None.	None

Source code in PytorchWildlife/utils/post_process.py

def save_detection_classification_timelapse_json(
    det_results, clf_results, output_path, det_categories=None, clf_categories=None,
    exclude_file_path=None, info={"detector": "megadetector_v5"}
):
    """
    Save detection and classification results to a JSON file in the specified format.

    Args:
        det_results (list):
            Detection results containing image ID, bounding boxes, detection category, and confidence.
        clf_results (list):
            Classification results containing image ID, classification category, and confidence.
        output_path (str):
            Path to save the output JSON file.
        det_categories (dict, optional):
            Dictionary of categories for detected objects. Defaults to None.
        clf_categories (dict, optional):
            Dictionary of categories for classified objects. Defaults to None.
        exclude_file_path (str, optional):
            We can exclude the some path sections from the image ID. Defaults to None.
    """
    json_results = {
        "info": info,
        "detection_categories": det_categories,
        "classification_categories": clf_categories,
        "images": []
    }

    for det_r in det_results:
        image_annotations = {
            "file": str(det_r["img_id"]).replace(exclude_file_path + os.sep, '') if exclude_file_path else str(det_r["img_id"]),
            "max_detection_conf": float(max(det_r["detections"].confidence)) if len(det_r["detections"].confidence) > 0 else '',
            "detections": []
        }

        for i in range(len(det_r["detections"])):
            det = det_r["detections"][i]
            normalized_bbox = [float(y) for y in det_r["normalized_coords"][i]]
            detection = {
                "category": str(det.class_id[0]),
                "conf": float(det.confidence[0]),
                "bbox": [normalized_bbox[0], normalized_bbox[1], normalized_bbox[2]-normalized_bbox[0], normalized_bbox[3]-normalized_bbox[1]],
                "classifications": []
            }

            # Find classifications for this detection
            for clf_r in clf_results:
                if clf_r["img_id"] == det_r["img_id"]:
                    detection["classifications"].append([str(clf_r["class_id"]), float(clf_r["confidence"])])

            image_annotations["detections"].append(detection)

        json_results["images"].append(image_annotations)

    with open(output_path, "w") as f:
        json.dump(json_results, f, indent=4)

save_detection_images(results, output_dir, input_dir=None, overwrite=False)

Save detected images with bounding boxes and labels annotated.

Parameters:

Name	Type	Description	Default
results	list or dict	Detection results containing image ID, detections, and labels.	required
output_dir	str	Directory to save the annotated images.	required
input_dir	str	Directory containing the input images. Default to None.	None
overwrite	bool	Whether overwriting existing image folders. Default to False.	False

Source code in PytorchWildlife/utils/post_process.py

def save_detection_images(results, output_dir, input_dir=None, overwrite=False):
    """
    Save detected images with bounding boxes and labels annotated.

    Args:
        results (list or dict):
            Detection results containing image ID, detections, and labels.
        output_dir (str):
            Directory to save the annotated images.
        input_dir (str):
            Directory containing the input images. Default to None.
        overwrite (bool):
            Whether overwriting existing image folders. Default to False.
    """
    box_annotator = sv.BoxAnnotator(thickness=4)
    lab_annotator = sv.LabelAnnotator(text_color=sv.Color.BLACK, text_thickness=4, text_scale=2)
    os.makedirs(output_dir, exist_ok=True)

    with sv.ImageSink(target_dir_path=output_dir, overwrite=overwrite) as sink: 
        if isinstance(results, list):
            for entry in results:
                annotated_img = lab_annotator.annotate(
                    scene=box_annotator.annotate(
                        scene=np.array(Image.open(entry["img_id"]).convert("RGB")),
                        detections=entry["detections"],
                    ),
                    detections=entry["detections"],
                    labels=entry["labels"],
                )
                if input_dir:
                    relative_path = os.path.relpath(entry["img_id"], input_dir)
                    save_path = os.path.join(output_dir, relative_path)
                    os.makedirs(os.path.dirname(save_path), exist_ok=True) 
                    image_name = relative_path 
                else:
                    image_name = os.path.basename(entry["img_id"])
                sink.save_image(
                    image=cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR), image_name=image_name
                )
        else:
            annotated_img = lab_annotator.annotate(
                scene=box_annotator.annotate(
                    scene=np.array(Image.open(results["img_id"]).convert("RGB")),
                    detections=results["detections"],
                ),
                detections=results["detections"],
                labels=results["labels"],
            )

            sink.save_image(
                image=cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR), image_name=os.path.basename(results["img_id"])
            )

save_detection_images_dots(results, output_dir, input_dir=None, overwrite=False)

Save detected images with bounding boxes and labels annotated.

Parameters:

Name	Type	Description	Default
results	list or dict	Detection results containing image ID, detections, and labels.	required
output_dir	str	Directory to save the annotated images.	required
input_dir	str	Directory containing the input images. Default to None.	None
overwrite	bool	Whether overwriting existing image folders. Default to False.	False

Source code in PytorchWildlife/utils/post_process.py

def save_detection_images_dots(results, output_dir, input_dir=None, overwrite=False):
    """
    Save detected images with bounding boxes and labels annotated.

    Args:
        results (list or dict):
            Detection results containing image ID, detections, and labels.
        output_dir (str):
            Directory to save the annotated images.
        input_dir (str):
            Directory containing the input images. Default to None.
        overwrite (bool):
            Whether overwriting existing image folders. Default to False.
    """
    dot_annotator = sv.DotAnnotator(radius=6)  
    lab_annotator = sv.LabelAnnotator(text_position=sv.Position.BOTTOM_RIGHT)   
    os.makedirs(output_dir, exist_ok=True)

    with sv.ImageSink(target_dir_path=output_dir, overwrite=overwrite) as sink:
        if isinstance(results, list):
            for i, entry in enumerate(results):
                if "img_id" in entry:
                    scene = np.array(Image.open(entry["img_id"]).convert("RGB"))
                    image_name = os.path.basename(entry["img_id"])
                else:
                    scene = entry["img"]
                    image_name = f"output_image_{i}.jpg" # default name if no image id is provided

                annotated_img = lab_annotator.annotate(
                    scene=dot_annotator.annotate(
                        scene=scene,
                        detections=entry["detections"],
                    ),
                    detections=entry["detections"],
                    labels=entry["labels"],
                )
                if input_dir:
                    relative_path = os.path.relpath(entry["img_id"], input_dir)
                    save_path = os.path.join(output_dir, relative_path)
                    os.makedirs(os.path.dirname(save_path), exist_ok=True) 
                    image_name = relative_path
                sink.save_image(
                    image=cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR), image_name=image_name
                )
        else:
            if "img_id" in results:
                scene = np.array(Image.open(results["img_id"]).convert("RGB"))
                image_name = os.path.basename(results["img_id"])
            else:
                scene = results["img"]
                image_name = "output_image.jpg" # default name if no image id is provided

            annotated_img = lab_annotator.annotate(
                scene=dot_annotator.annotate(
                    scene=scene,
                    detections=results["detections"],
                ),
                detections=results["detections"],
                labels=results["labels"],
            )   
            sink.save_image(
                image=cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR), image_name=image_name
            )

save_detection_json(det_results, output_dir, categories=None, exclude_category_ids=[], exclude_file_path=None)

Save detection results to a JSON file.

Parameters:

Name	Type	Description	Default
det_results	list	Detection results containing image ID, bounding boxes, category, and confidence.	required
output_dir	str	Path to save the output JSON file.	required
categories	list	List of categories for detected objects. Defaults to None.	None
exclude_category_ids	list	List of category IDs to exclude from the output. Defaults to []. Category IDs can be found in the definition of each models.	[]
exclude_file_path	str	We can exclude the some path sections from the image ID. Defaults to None.	None

Source code in PytorchWildlife/utils/post_process.py

def save_detection_json(det_results, output_dir, categories=None, exclude_category_ids=[], exclude_file_path=None):
    """
    Save detection results to a JSON file.

    Args:
        det_results (list):
            Detection results containing image ID, bounding boxes, category, and confidence.
        output_dir (str):
            Path to save the output JSON file.
        categories (list, optional):
            List of categories for detected objects. Defaults to None.
        exclude_category_ids (list, optional):
            List of category IDs to exclude from the output. Defaults to []. Category IDs can be found in the definition of each models.
        exclude_file_path (str, optional):
            We can exclude the some path sections from the image ID. Defaults to None.
    """
    json_results = {"annotations": [], "categories": categories}

    for det_r in det_results:

        # Category filtering
        img_id = det_r["img_id"]
        category = det_r["detections"].class_id

        bbox = det_r["detections"].xyxy.astype(int)[~np.isin(category, exclude_category_ids)]
        confidence =  det_r["detections"].confidence[~np.isin(category, exclude_category_ids)]
        category = category[~np.isin(category, exclude_category_ids)]

        # if not all([x in exclude_category_ids for x in category]):
        json_results["annotations"].append(
            {
                "img_id": img_id.replace(exclude_file_path + os.sep, '') if exclude_file_path else img_id,
                "bbox": bbox.tolist(),
                "category": category.tolist(),
                "confidence": confidence.tolist(),
            }
        )

    with open(output_dir, "w") as f:
        json.dump(json_results, f, indent=4)

save_detection_json_as_dots(det_results, output_dir, categories=None, exclude_category_ids=[], exclude_file_path=None)

Save detection results to a JSON file in dots format.

Parameters:

Name	Type	Description	Default
det_results	list	Detection results containing image ID, bounding boxes, category, and confidence.	required
output_dir	str	Path to save the output JSON file.	required
categories	list	List of categories for detected objects. Defaults to None.	None
exclude_category_ids	list	List of category IDs to exclude from the output. Defaults to []. Category IDs can be found in the definition of each models.	[]
exclude_file_path	str	We can exclude the some path sections from the image ID. Defaults to None.	None

Source code in PytorchWildlife/utils/post_process.py

def save_detection_json_as_dots(det_results, output_dir, categories=None, exclude_category_ids=[], exclude_file_path=None):
    """
    Save detection results to a JSON file in dots format.

    Args:
        det_results (list):
            Detection results containing image ID, bounding boxes, category, and confidence.
        output_dir (str):
            Path to save the output JSON file.
        categories (list, optional):
            List of categories for detected objects. Defaults to None.
        exclude_category_ids (list, optional):
            List of category IDs to exclude from the output. Defaults to []. Category IDs can be found in the definition of each models.
        exclude_file_path (str, optional):
            We can exclude the some path sections from the image ID. Defaults to None.
    """
    json_results = {"annotations": [], "categories": categories}

    for det_r in det_results:

        # Category filtering
        img_id = det_r["img_id"]
        category = det_r["detections"].class_id

        bbox = det_r["detections"].xyxy.astype(int)[~np.isin(category, exclude_category_ids)]
        dot = np.array([[np.mean(row[::2]), np.mean(row[1::2])] for row in bbox])
        confidence =  det_r["detections"].confidence[~np.isin(category, exclude_category_ids)]
        category = category[~np.isin(category, exclude_category_ids)]

        # if not all([x in exclude_category_ids for x in category]):
        json_results["annotations"].append(
            {
                "img_id": img_id.replace(exclude_file_path + os.sep, '') if exclude_file_path else img_id,
                "dot": dot.tolist(),
                "category": category.tolist(),
                "confidence": confidence.tolist(),
            }
        )

    with open(output_dir, "w") as f:
        json.dump(json_results, f, indent=4)

save_detection_timelapse_json(det_results, output_dir, categories=None, exclude_category_ids=[], exclude_file_path=None, info={'detector': 'megadetector_v5'})

Save detection results to a JSON file.

Parameters:

Name	Type	Description	Default
det_results	list	Detection results containing image ID, bounding boxes, category, and confidence.	required
output_dir	str	Path to save the output JSON file.	required
categories	list	List of categories for detected objects. Defaults to None.	None
exclude_category_ids	list	List of category IDs to exclude from the output. Defaults to []. Category IDs can be found in the definition of each models.	[]
exclude_file_path	str	Some time, Timelapse has path issues. We can exclude the some path sections from the image ID. Defaults to None.	None
info	dict	Default Timelapse info. Defaults to {"detector": "megadetector_v5}.	{'detector': 'megadetector_v5'}

Source code in PytorchWildlife/utils/post_process.py

def save_detection_timelapse_json(
    det_results, output_dir, categories=None,
    exclude_category_ids=[], exclude_file_path=None, info={"detector": "megadetector_v5"}
    ):
    """
    Save detection results to a JSON file.

    Args:
        det_results (list):
            Detection results containing image ID, bounding boxes, category, and confidence.
        output_dir (str):
            Path to save the output JSON file.
        categories (list, optional):
            List of categories for detected objects. Defaults to None.
        exclude_category_ids (list, optional):
            List of category IDs to exclude from the output. Defaults to []. Category IDs can be found in the definition of each models.
        exclude_file_path (str, optional):
            Some time, Timelapse has path issues. We can exclude the some path sections from the image ID. Defaults to None.
        info (dict, optional):
            Default Timelapse info. Defaults to {"detector": "megadetector_v5}.
    """

    json_results = {
        "info": info,
        "detection_categories": categories,
        "images": []
    }

    for det_r in det_results:

        img_id = det_r["img_id"]
        category_id_list = det_r["detections"].class_id

        bbox_list = det_r["detections"].xyxy.astype(int)[~np.isin(category_id_list, exclude_category_ids)]
        confidence_list =  det_r["detections"].confidence[~np.isin(category_id_list, exclude_category_ids)]
        normalized_bbox_list = np.array(det_r["normalized_coords"])[~np.isin(category_id_list, exclude_category_ids)]
        category_id_list = category_id_list[~np.isin(category_id_list, exclude_category_ids)]

        # if not all([x in exclude_category_ids for x in category_id_list]):
        image_annotations = {
            "file": img_id.replace(exclude_file_path + os.sep, '') if exclude_file_path else img_id,
            "max_detection_conf": float(max(confidence_list)) if len(confidence_list) > 0 else '',
            "detections": []
        }
        for i in range(len(bbox_list)):
            normalized_bbox = [float(y) for y in normalized_bbox_list[i]]
            detection = {
                "category": str(category_id_list[i]),
                "conf": float(confidence_list[i]),
                "bbox": [normalized_bbox[0], normalized_bbox[1], normalized_bbox[2]-normalized_bbox[0], normalized_bbox[3]-normalized_bbox[1]],
                "classifications": []
            }

            image_annotations["detections"].append(detection)

        json_results["images"].append(image_annotations)

    with open(output_dir, "w") as f:
        json.dump(json_results, f, indent=4)