app.py

#!/usr/bin/env python3
"""
PiClassifier - A simple PyQt5 application to classify Raspberry Pi objects from a camera stream.
"""

import sys
from typing import Optional, Tuple, Any
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QHBoxLayout, QWidget, QGraphicsView, QGraphicsScene, QGraphicsRectItem
from PyQt5 import QtCore
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QFont, QPalette, QPen, QBrush, QColor
from PyQt5.QtWidgets import QGraphicsTextItem

import numpy as np

from picamera2 import Picamera2
from picamera2.previews.qt import QGlPicamera2
from picamera2.devices import IMX500
from picamera2.devices.imx500 import NetworkIntrinsics
from picamera2.devices.imx500.postprocess import softmax
from picamera2.job import Job


class PiClassifier(QMainWindow):
    """
    Main PiClassifier application window.

    This class creates a PyQt5-based GUI application for real-time image classification
    using a Raspberry Pi camera and IMX500 neural network processor. The application
    displays live camera feed, classification results, confidence scores, and a
    histogram showing all class probabilities.

    Attributes:
        picam2: Picamera2 instance for camera control
        imx500: IMX500 neural network processor instance
        intrinsics: NetworkIntrinsics object containing classification parameters
        camera_widget: QGlPicamera2 widget for camera preview
        histogram_widget: QGraphicsView widget for probability histogram
        histogram_scene: QGraphicsScene containing histogram elements
        label: QLabel displaying current classification result
        confidence_label: QLabel displaying confidence percentage
        last_detection: Tuple storing the last detection result and all probabilities
    """

    def __init__(self) -> None:
        """
        Initialize the PiClassifier application.

        Sets up the camera, initializes the UI components, and starts the camera
        for real-time classification.
        """
        super().__init__()
        self.init_camera()
        self.init_ui()
        self.start_camera()

    def init_ui(self) -> None:
        """
        Initialize the user interface.

        Creates and arranges all GUI components including:
        - Top section with detection and confidence labels
        - Camera preview widget
        - Histogram widget for probability visualization

        The layout uses fixed positioning for consistent placement of elements.
        """
        self.setWindowTitle("PiClassifier")
        self.setGeometry(100, 100, 800, 600)

        # Create central widget and layout
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout()
        central_widget.setLayout(layout)

        # Create top section container
        top_widget = QWidget()
        top_widget.setFixedHeight(50)  # Fixed height for top section

        # Create a label at the top
        self.label = QLabel("")
        self.label.setParent(top_widget)
        self.label.setFixedSize(300, 30)  # Set explicit size

        # Set font for the label
        font = QFont()
        font.setPointSize(24)
        font.setBold(True)
        self.label.setFont(font)

        # Create "Detected:" title
        detected_label = QLabel("Detected:")
        detected_label.setParent(top_widget)
        detected_font = QFont()
        detected_font.setPointSize(16)
        detected_label.setFont(detected_font)

        # Create "Confidence:" title
        confidence_title_label = QLabel("Confidence:")
        confidence_title_label.setParent(top_widget)
        confidence_title_font = QFont()
        confidence_title_font.setPointSize(16)
        confidence_title_label.setFont(confidence_title_font)

        # Create confidence value label
        self.confidence_label = QLabel("")
        self.confidence_label.setParent(top_widget)
        self.confidence_label.setFixedSize(100, 30)  # Set explicit size
        confidence_font = QFont()
        confidence_font.setPointSize(24)
        confidence_font.setBold(True)
        self.confidence_label.setFont(confidence_font)

        # Position widgets with fixed coordinates
        detected_label.move(10, 15)
        self.label.move(110, 10)  # Moved right to avoid overlap
        confidence_title_label.move(400, 15)  # Fixed position at center
        self.confidence_label.move(520, 10)  # Moved right by same amount

        # Add top widget to main layout
        layout.addWidget(top_widget)

        # Create camera preview widget
        bg_colour = self.palette().color(QPalette.Background).getRgb()[:3]
        self.camera_widget = QGlPicamera2(self.picam2, width=800, height=600, bg_colour=bg_colour)
        self.camera_widget.setFixedSize(800, 600)  # Explicitly set the size
        layout.addWidget(self.camera_widget)  # Add without stretch factor

        # Create histogram widget
        self.histogram_widget = QGraphicsView()
        self.histogram_scene = QGraphicsScene()
        self.histogram_widget.setScene(self.histogram_scene)
        self.histogram_widget.setFixedHeight(200)
        self.histogram_widget.setStyleSheet("background-color: #808080;")  # Medium grey background
        self.histogram_widget.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.histogram_widget.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        layout.addWidget(self.histogram_widget)
        self.camera_widget.done_signal.connect(self.camera_callback, type=QtCore.Qt.QueuedConnection)

    def init_camera(self) -> None:
        """
        Initialize the camera configuration.

        Sets up the Picamera2 instance with preview configuration including:
        - Main stream: 800x600 YUV420 format
        - Sensor output: 2028x1520 resolution
        - Frame rate: 25 FPS

        The camera is configured but not started until start_camera() is called.
        """
        self.picam2 = Picamera2()
        main = {'size': (800, 600), 'format': 'YUV420'}
        sensor = {'output_size': (2028, 1520)}
        controls = {'FrameRate': 25}
        config = self.picam2.create_preview_configuration(main, sensor=sensor, controls=controls)
        self.picam2.configure(config)

    def start_camera(self) -> None:
        """
        Start the camera and neural network processing.

        Initializes the IMX500 neural network processor, loads classification labels,
        and begins real-time camera capture with metadata processing. Sets up the
        continuous capture loop for classification.

        Raises:
            FileNotFoundError: If network.rpk or labels.txt files are not found
        """
        self.imx500 = IMX500("network.rpk")

        self.intrinsics = NetworkIntrinsics()
        self.intrinsics.task = "classification"
        with open("labels.txt") as f:
            self.intrinsics.labels = f.read().splitlines()
        self.intrinsics.update_with_defaults()

        self.imx500.show_network_fw_progress_bar()
        self.picam2.start()

        self.last_detection = None

        # And fire off the first asynchronous request for the next frame's metadata.
        self.picam2.capture_metadata(signal_function=self.camera_widget.signal_done)

    def camera_callback(self, job: Job) -> None:
        """
        Handle camera frame processing and update UI with classification results.

        This method is called for each captured frame. It processes the metadata
        to extract classification results, updates the UI labels with the top
        detection and confidence, and refreshes the histogram with all class
        probabilities.

        Args:
            job: Camera job object containing frame metadata
        """
        metadata = job.get_result()
        # Fire off the asynchronous request for the next frame's metadata right away.
        self.picam2.capture_metadata(signal_function=self.camera_widget.signal_done)

        top_detection, all_probs = self.parse_detection(metadata)
        if top_detection is not None:
            detected = self.intrinsics.labels[top_detection[0]]
            detected = "" if detected == "(None)" else detected
            label = f"{detected}"
            confidence = f"{int(top_detection[1] * 100)}%"
            self.label.setText(label)
            self.confidence_label.setText(confidence)

            # Update histogram
            self.update_histogram(all_probs)

    def parse_detection(self, metadata: dict) -> Optional[Tuple[Tuple[int, float], np.ndarray]]:
        """
        Parse neural network output from camera metadata.

        Extracts classification probabilities from the IMX500 neural network
        output. If the network hasn't finished processing the current frame,
        returns the last known detection result.

        Args:
            metadata: Camera metadata containing neural network outputs

        Returns:
            Optional tuple containing (top_detection, all_probabilities) where:
                - top_detection: (class_index, confidence) of highest probability class
                - all_probabilities: numpy array of all class probabilities
            Returns None if no valid detection is available
        """
        np_outputs = self.imx500.get_outputs(metadata)
        if np_outputs is None:
            # Occasionally, the IMX500 may not have finished processing this frame, in which
            # case we should re-use the last result.
            return self.last_detection

        np_output = np_outputs[0]
        if self.intrinsics.softmax:  # this network recommends applying softmax
            np_output = softmax(np_output)

        top_index = np.argmax(np_output)  # get the index of the highest confidence class
        self.last_detection = (top_index, np_output[top_index]), np_output
        return self.last_detection

    def update_histogram(self, probabilities: Optional[np.ndarray]) -> None:
        """
        Update the histogram with current classification probabilities.

        Creates a vertical bar chart showing the probability distribution across
        all classification classes. Each bar represents one class with height
        proportional to its probability. Bars are colored with a blue gradient
        based on probability intensity.

        Args:
            probabilities: numpy array of classification probabilities for all classes
        """
        # Clear previous items
        self.histogram_scene.clear()

        if probabilities is None or len(probabilities) == 0:
            return

        # Calculate dimensions
        num_labels = len(self.intrinsics.labels)
        bar_width = 800 / num_labels  # Total width divided by number of bars
        max_height = 150  # Maximum bar height

        # Create bars for each label
        for i, (label, prob) in enumerate(zip(self.intrinsics.labels, probabilities)):
            # Calculate bar dimensions
            bar_height = prob * max_height
            x = i * bar_width
            y = max_height - bar_height  # Flip Y coordinate (Qt has origin at top)

            # Create bar rectangle
            bar = QGraphicsRectItem(x, y, bar_width - 2, bar_height)  # -2 for spacing

            # Set bar color (blue gradient based on probability)
            color_intensity = int(255 * prob)
            bar_color = QColor(0, 100, color_intensity)
            bar.setBrush(QBrush(bar_color))
            bar.setPen(QPen(QColor(0, 0, 0), 1))

            # Add bar to scene
            self.histogram_scene.addItem(bar)

            # Add label text below the bar
            text_item = QGraphicsTextItem(label)
            text_item.setPos(x + bar_width/2 - text_item.boundingRect().width()/2, max_height + 5)
            self.histogram_scene.addItem(text_item)

            # Add probability percentage above the bar
            prob_text = QGraphicsTextItem(f"{int(prob * 100)}%")
            prob_text.setPos(x + bar_width/2 - prob_text.boundingRect().width()/2, y - 20)
            self.histogram_scene.addItem(prob_text)


def main() -> None:
    """
    Main application entry point.

    Creates and configures the QApplication, initializes the PiClassifier
    main window, and starts the Qt event loop. This function handles the
    complete application lifecycle.
    """
    app = QApplication(sys.argv)

    # Set application properties
    app.setApplicationName("PiClassifier")
    app.setApplicationVersion("1.0")
    app.setOrganizationName("Pi Classifier Project")

    # Create and show main window
    window = PiClassifier()
    window.show()

    # Start the application event loop
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()