mail@pastecode.io avatar
a month ago
3.2 kB
import tkinter as tk
from tkinter import Canvas, Label
from PIL import Image, ImageTk
from ultralytics import YOLO
import cv2
import math

# Function to update the canvas and label with the latest image and information
def update_image():
    success, img = cap.read()
    results = model(img, stream=True)

    for r in results:
        boxes = r.boxes

        for box in boxes:
            x1, y1, x2, y2 = box.xyxy[0]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)

            # Draw bounding box on the canvas
            canvas.create_rectangle(x1, y1, x2, y2, outline="pink", width=2)

            confidence = math.ceil((box.conf[0] * 100)) / 100
            class_name = classNames[int(box.cls[0])]

            # Display confidence and class name on the canvas
            canvas.create_text(x1, y1 - 10, text=f"Confidence: {confidence}", fill="red", anchor=tk.SW)
            canvas.create_text(x1, y1 - 30, text=f"Class: {class_name}", fill="blue", anchor=tk.SW)

            # Update label with confidence and class name
            info_label.config(text=f"Confidence: {confidence}, Class: {class_name}")

    # Convert image to PhotoImage format
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(img)
    img = ImageTk.PhotoImage(img)

    # Update the image on the canvas
    canvas.imgtk = img
    canvas.create_image(0, 0, anchor=tk.NW, image=img)

    # Schedule the update after 10 milliseconds
    root.after(10, update_image)

# Create the Tkinter window
root = tk.Tk()
root.title("Object Detection with YOLO")

# Set up the canvas for displaying the webcam feed
canvas = Canvas(root, width=640, height=480)

# Set up label for displaying information under the webcam stream
info_label = Label(root, text="", font=("Helvetica", 12))

# Start webcam
cap = cv2.VideoCapture(0)
cap.set(3, 640)
cap.set(4, 480)

# YOLO model
model = YOLO("yolov8n.pt")

# Object classes
classNames = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
              "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
              "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
              "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
              "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
              "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
              "carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
              "diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
              "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
              "teddy bear", "hair drier", "toothbrush"

# Schedule the initial update
root.after(10, update_image)

# Run the Tkinter main loop

# Release webcam and close OpenCV windows
Leave a Comment