pytorch_inf.py

from ultralytics import YOLO
import pandas as pd
from habitat_classification import classify_habitat, get_inconclusive
import time
import cv2

def get_habitat_data(class_names, species_data):
    if not class_names:
        inconclusive_res = get_inconclusive()
        return inconclusive_res[0], inconclusive_res[1]
    else:
        habitat_res = classify_habitat(class_names, species_data)
        return habitat_res[0], habitat_res[1]


def display_detections(frame, results):
    detections = []
    annotations = []
    for result in results:
        if result.boxes.data.nelement() != 0:
            for detection in result.boxes.data:
                x1, y1, x2, y2, conf, cls = detection.cpu().numpy()
                class_name = result.names[int(cls)]
                detections.append([class_name, conf])
                annotations.append(((int(x1), int(y1), int(x2), int(y2)), class_name))
            outframe = result.plot()

            return frame, outframe, annotations, pd.DataFrame(detections, columns=["Class", "Confidence"])

    print("No detections")
    return frame, None, [], pd.DataFrame(columns=["Class", "Confidence"])


def process_image(input_image, model):
    corrected_image_rgb = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
    results = model(corrected_image_rgb, stream=True)
    return display_detections(corrected_image_rgb, results)


def run_pytorch_inference(input_image, pt_model, species_data):
    model = YOLO(pt_model)
    model.conf = 0.2
    start_time = time.time()
    frame, outframe, annotations, detections = process_image(input_image, model)
    if detections.empty:
        habitat, probabilities = get_inconclusive()
        detections = pd.DataFrame(columns=["Class", "Confidence"])
    else:
        habitat, probabilities = get_habitat_data(detections['Class'].tolist(),species_data)
    end_time = time.time()
    duration = end_time - start_time
    return f"Executed in {duration:.2f}s", (frame, annotations), habitat, probabilities, outframe, detections