Spaces:
Sleeping
Sleeping
File size: 2,867 Bytes
58ddda3 911c0ac 58ddda3 911c0ac 58ddda3 911c0ac 631ca76 911c0ac d304ca0 911c0ac 631ca76 911c0ac 631ca76 911c0ac 631ca76 911c0ac 631ca76 911c0ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
from datetime import datetime
from data_models.bbox_manager import BoundingBoxManager, add_bboxes_to_db
from data_models.image_manager import ImageManager
from ultralytics import YOLO
model = None
bbox_manager = BoundingBoxManager()
image_manager = ImageManager()
def add_image_to_db(result, image_manager, park_id):
"""
Adds an image to the database.
Args:
result: YOLO result object containing the image details.
image_manager (ImageManager): Instance of ImageManager to interact with the database.
park_id (int): The ID of the associated park.
Returns:
int: The ID of the added image.
"""
image_name = result.path.split("/")[-1]
current_datetime = datetime.now()
created_at = current_datetime.strftime("%Y-%m-%d %H:%M:%S")
try:
print(f"Adding image '{image_name}' to the database...")
image = image_manager.add_image(
name=image_name, created_at=created_at, park_id=park_id
)
image_id = image_manager.get_image_id(image_name)
print(f"Image '{image_name}' added with ID: {image_id}")
return image_id
except Exception as e:
print(f"Error adding image '{image_name}' to the database: {e}")
return None
def load_model(model_file):
"""
Load the YOLO model.
Args:
model_file (str): Path to the YOLO model.
Returns:
YOLO: The YOLO model.
"""
global model
if model:
return model
try:
print(f"Loading YOLO model from file: {model_file}...")
model = YOLO(f"models/{model_file}")
print(f"Model loaded successfully!")
return model
except Exception as e:
print(f"Error loading YOLO model: {e}")
return None
def process_YOLO_predictions(park_id: int, image: str, model_file: str):
"""
Process predictions for a list of image folders, adding images and bounding boxes to the database using the YOLO model.
Args:
list_folder (list): List of image folders.
model_file (str): Path to the YOLO model.
Raises:
Exception: If an error occurs during processing.
"""
print(f"Initializing processing of predictions...")
model = load_model(model_file)
if not model:
raise Exception("Error loading YOLO model. Aborting processing.")
predictions = model(image, stream=True, show=False)
for result in predictions:
print(f"Processing result for image: {result.path.split('/')[-1]}...")
image_id = add_image_to_db(result, image_manager, park_id)
if image_id:
print(f"Adding bounding boxes to the database...")
add_bboxes_to_db(result, bbox_manager, image_id)
print(f"Processing complete. Closing database connections.")
bbox_manager.close_connection()
image_manager.close_connection()
|