File size: 3,048 Bytes
a913145 1f2862c a913145 9071273 a913145 1f2862c a913145 1f2862c a913145 |
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 97 98 |
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import os
import shutil
import cv2
import subprocess
from pathlib import Path
# Method to detect object counts with given confidence score
# Helper function to get key by value from a dictionary
def getKeyByValue(dictionary, value):
for key, val in dictionary.items():
if val == value:
return key
return None # Return None if value not found
# Usage in YOLO prediction
#class_name = "car" # The class you want to filter, e.g., "car"
def detectObjectsAndCount(imageFile, confidence_score, class_type):
# Load a pre-trained YOLOv8 model (e.g., YOLOv8n)
model = YOLO("yolov8n.pt") # Use other variants like yolov8s.pt for better accuracy
# Define a custom directory to save the results
custom_save_dir = "/app/runs/detect/predict"
custom_read_dir = "./"
class_index = getKeyByValue(model.names, class_type)
count_of_class_type = 0;
# Run inference on an image or video
results = model.predict(
source=f"{imageFile}",
save=True,
conf=confidence_score,
save_dir=custom_save_dir, # Specify the custom output directory
exist_ok=True, # Prevent creating new subdirectories
classes=[class_index]
)
# Count the number of objects of the specified class
for result in results:
for box in result.boxes:
class_name = model.names[int(box.cls)]
if class_name == class_type:
count_of_class_type += 1
# Path to the saved image in the custom directory
image_path = f"{custom_save_dir}/{imageFile}"
pred_dir = Path("/app/runs/detect/predict")
stem = Path(image_path).stem # IMG_2451
matches = list(pred_dir.glob(f"{stem}.*"))
if not matches:
raise FileNotFoundError(f"No output image found for {stem} in {pred_dir}")
# Load and display the image
#img = mpimg.imread(str(matches[0]))
# Open the image using Pillow
img = Image.open(str(matches[0]))
# Prepare to draw text on the image
draw = ImageDraw.Draw(img)
font_path = "/System/Library/Fonts/Supplemental/Arial.ttf" # Update this path if necessary
font_size = 30 # Specify the desired font size
try:
font = ImageFont.truetype(font_path, font_size)
except OSError:
font = ImageFont.load_default() # You can load a custom font if needed
# Text to display
text = f"Number of {class_type}s are: {count_of_class_type}"
# Define text position (at the top of the image)
text_position = (200, 10)
# Define text color (dark blue)
text_color = (0, 0, 139) # RGB value for dark blue
# Add text to the image
draw.text(text_position, text, fill=text_color, font=font)
# Display the modified image with the overlay
#plt.figure(figsize=(8, 12))
#plt.imshow(img)
#plt.axis("off") # Hide axes
#plt.show()
return img, count_of_class_type
|