birdfsd_yolov5.prediction package

Submodules

birdfsd_yolov5.prediction.auto_prediction_schedule module

birdfsd_yolov5.prediction.auto_prediction_schedule.auto_prediction_pipeline(opts_file: Optional[str] = None, show_opts: bool = False) None[source]

A pipeline to run the prediction module.

A prediction pipeline that is ntended to be used as systemctl service or inside a GitHub actions workflow.

Returns

None

birdfsd_yolov5.prediction.predict module

class birdfsd_yolov5.prediction.predict.LoadModel(weights: str, model_version: str)[source]

Bases: object

This class is used to load a model from a given path.

model() Module[source]

Loads a pretrained YOLOv5 model for a given dataset.

The model is loaded using the torch.hub.load function.

Returns

a YOLOv5 model.

Return type

torch.nn.Module

class birdfsd_yolov5.prediction.predict.Predict(weights: str, model_version: str, project_ids: Optional[str] = None, tasks_range: Optional[str] = None, predict_all: bool = False, one_task: Optional[int] = None, multithreading: bool = True, delete_if_no_predictions: bool = True, if_empty_apply_label: Optional[str] = None, get_tasks_with_api: bool = False, verbose: bool = False)[source]

Bases: LoadModel, _Headers

Prediction and preprocessing class.

This class is used to predict bounding boxes and classes for images in a given project. It uses the YOLOv5 model to predict bounding boxes and then posts the predictions to the Label Studio server.

headers

Headers for the requests.

Type

dict

model

YOLOv5 model.

Type

torch.nn.Module

model_version

Model version to use.

Type

str

project_ids

Ids of the projects to predict.

Type

Optional[str]

tasks_range

Range of tasks to predict.

Type

Optional[str]

predict_all

Predict all tasks in the project(s).

Type

bool

one_task

Predict a single task.

Type

Optional[int]

multithreading

Use multithreading.

Type

bool

delete_if_no_predictions

Delete tasks that have no predictions.

Type

bool

if_empty_apply_label

Apply a label to tasks that have no predictions.

Type

str

get_tasks_with_api

Get tasks with label-studio API instead of MongoDB.

Type

bool

db

An instance of mongoDB client (if connection string exists in .env).

Type

pymongo.database.Database

flush

Used to flush temp files written to disk during prediction.

Type

list

apply_predictions() None[source]

This function applies predictions to label studio tasks.

Returns

None

download_image(img_url: str) str[source]

Download an image from a URL and save it to a local file.

Parameters

img_url (str) – The URL of the image to download.

Returns

Path to the temporary image file.

Return type

str

get_all_tasks(project_ids: List[str]) list[source]

Fetch all tasks from the project.

This function fetches all tasks from the project.

Parameters
  • project_ids (str) – Comma-seperated string of project ids.

  • project_ids – List[str]:

Returns

A list of tasks.

Return type

list

get_task(_task_id: int) dict[source]

This function returns a single task from a project.

Parameters

_task_id (int) – The id of the task to be returned.

Returns

A dictionary containing the task data.

Return type

dict

post_prediction(task: dict) Optional[Union[str, dict]][source]

Prepares the POST request data of the prediction results.

This function is called by the predict method. It takes a task as an argument and performs the following steps:

  1. It downloads the image from the task’s data field.

  2. It runs the image through the model and gets the predictions.

  3. It converts the predictions to the format required by Label Studio.

  4. It posts the predictions to Label Studio.

If the task has no data, it skips the task.

If the task has no predictions, it deletes the task if delete_if_no_predictions is set to True.

If if_empty_apply_label is set to a label, it applies the string of if_empty_apply_label if not set to None.

Parameters

task (dict) – A dictionary with the task data.

Returns

_skipped or task dictionary.

Return type

Optional[Union[str, dict]]

Raises
  • UnidentifiedImageError – If image is not identified.

  • OSError – If image is not downloaded.

  • Exception – If any other exception occurs.

pred_exists(task: dict) Optional[bool][source]

Check if a prediction already exists with the current model version.

Parameters

task – A dictionary with the prediction’s data.

Returns

True if a prediction with the current model version exists. Otherwise, None.

Return type

Optional[bool]

pred_post(results: list, scores: list, task_id: int) dict[source]

Creates a dictionary for a prediction POST data.

This function is used to create an API POST request of a single prediction results.

Parameters
  • results (list) – the results of the model for all predictions in a single image.

  • scores (list) – the scores of all detections predicted by the model.

  • task_id (int) – the task id.

Returns

The prediction results.

Return type

dict

static pred_result(x: float, y: float, w: float, h: float, score: float, label: str) dict[source]

Creates a dictionary for an individual prediction result.

This function takes in the x, y, width, height, score, and label of a prediction and returns a dictionary with the prediction’s information.

Parameters
  • x (float) – The x coordinate of the prediction.

  • y (float) – The y coordinate of the prediction.

  • w (float) – The width of the prediction.

  • h (float) – The height of the prediction.

  • score (float) – The confidence score of the prediction.

  • label (str) – The label of the prediction.

Returns

A dictionary with the prediction’s data.

Return type

dict

static selected_tasks(tasks: list, start: int, end: int) list[source]

Selects label-studio tasks in a user-specified range

Creates a list of tasks from the given list of tasks, whose id is in the range [start, end].

Parameters
  • tasks (list) – A list of tasks id.

  • start (int) – The start of the range.

  • end (int) – The end of the range.

Returns

A list of tasks.

Return type

list

single_task(task_id: int) list[source]

Get a single task by its id.

Parameters

task_id (int) – The id of the task to get.

Returns

A list containing the task data.

Return type

list

yolo_to_ls(x: float, y: float, width: float, height: float, score: float, n: int) tuple[source]

Converts YOLOv5 output to a tuple of the form: (x, y, w, h, conf, n)

Parameters
  • x (float) – The x coordinate of the center of the bounding box.

  • y (float) – The y coordinate of the center of the bounding box.

  • width (float) – The width of the bounding box.

  • height (float) – The height of the bounding box.

  • score (float) – The confidence score of the bounding box.

  • n (int) – The index of the class label.

Returns

(x, y, width, height, score, label)

Return type

tuple

Module contents