MLBench/ports-classification

1import gradio as gr2from ultralytics import YOLO3from PIL import Image4 5# --- Documentation Strings ---6 7USAGE_GUIDELINES = """8## 1. Quick Start Guide: Processing an Image9This application uses a specialized YOLO model for object detection to automatically locate and classify specific "Port" features within an image.10 111.  **Upload Image**: Click the 'Upload Image' box and select your image file (JPG or PNG).122.  **Run**: Click the **"Process Image"** button.133.  **Review**: The 'Output Image with labels' will display the detected ports, marked by bounding boxes, class labels, and confidence scores.14"""15 16INPUT_EXPLANATION = """17## 2. Expected Inputs and Best Practices18 19| Input Field | Purpose | Requirement |20| :--- | :--- | :--- |21| **Upload Image** | The image containing the ports or features you wish to classify. | Must be a standard image file (JPG, PNG). |22 23### Optimal Image Conditions24*   **Clarity and Focus:** Ensure the ports are clearly visible and in focus. Blurry images significantly reduce detection accuracy.25*   **Lighting:** Use well-lit images. Shadows or low-light conditions can cause the model to miss detections or misclassify objects.26*   **Framing:** The ports should occupy a significant portion of the image. Highly zoomed-out images may treat ports as background noise.27"""28 29OUTPUT_EXPLANATION = """30## 3. Expected Output 31 32The output image features graphical annotations provided by the AI model. Each annotation consists of three key components:33 34| Annotation Component | Description | Importance |35| :--- | :--- | :--- |36| **Bounding Box** | A colored rectangle drawn tightly around the detected object (the port). | Visually indicates the exact location of the object. |37| **Class Label** | The name of the predicted port type (e.g., USB-A, HDMI, Ethernet). | Represents the model's classification of the object within the box. |38| **Confidence Score** | A percentage (e.g., 0.92) representing the model's certainty in its prediction. | Scores below 0.50 (50%) should typically be treated as unreliable detections. |39 40"""41 42TECHNICAL_TIPS = """43## 4. Troubleshooting and Tips44 45*   **No Detections:** If the output image is identical to the input with no boxes, the model either failed to find any objects or the confidence for all potential detections was too low to be displayed. Try improving the lighting or clarity of the input photo.46*   **Testing with Examples:** Use the provided example images to quickly verify the application's functionality and see what types of objects the model is trained to recognize.47*   **Processing Time:** Processing time depends on the image resolution and the complexity of the scene. Large images require more computational resources for detection.48"""49 50# --------------------51# Core Pipeline Functions (Kept AS IS)52# --------------------53 54# Load the YOLO model weights55# Ensure 'best.pt' is available in the run directory56try:57    model = YOLO("./best.pt")58except Exception as e:59    print(f"Error loading model weights: {e}. Using a placeholder for demonstration.")60    # Define a placeholder function if model loading fails61    def placeholder_predict(img):62        return Image.new('RGB', img.size, color='red')63    model = placeholder_predict64 65 66def process_img(img: Image.Image):67    if img is None:68        gr.Warning("Please upload an image before processing.")69        return None70        71 72    # Perform prediction73    result = model.predict(img)74    75    # r.plot() returns a BGR numpy array76    for r in result:77        im_bgr = r.plot()78        # Convert BGR (OpenCV standard) to RGB (PIL/Gradio standard)79        return Image.fromarray(im_bgr[..., ::-1])80 81 82# --------------------83# Gradio UI84# --------------------85 86with gr.Blocks(title="Port Classification App") as demo:87    gr.Markdown("<h1 style='text-align: center;'> Port Classification and Object Detection </h1>")88    89    # 1. Guidelines Accordion90    with gr.Accordion("Tips & User Guidelines", open=False):91        gr.Markdown(USAGE_GUIDELINES)92        gr.Markdown("---")93        gr.Markdown(INPUT_EXPLANATION)94        gr.Markdown("---")95        gr.Markdown(OUTPUT_EXPLANATION)96        gr.Markdown("---")97        gr.Markdown(TECHNICAL_TIPS)98 99    # 2. Interface Definition100    with gr.Row():101        with gr.Column():102            gr.Markdown("## Step 1: Upload Port Image ")103            upload_img = gr.Image(label=" Upload Image", type="pil")104            gr.Markdown("## Step 2: Click Process Image ")105            classify_img_button = gr.Button(value=" Process Image", variant="primary")106 107        with gr.Column():108            gr.Markdown("## Result ")109            output_img = gr.Image(label=" Output Image with Detected Ports")110            111    # 3. Examples Section112    gr.Markdown("## Examples ")113    with gr.Row():114        gr.Examples(115            examples=[116                ["./examples/01.jpg"],117                ["./examples/02.jpg"],118                ["./examples/03.jpg"],119                ["./examples/04.jpg"],120                ["./examples/05.jpg"],121                ["./examples/06.jpg"],122            ],123            inputs=upload_img,124            label=" Click to load and test an example image"125        )126        127    # Event Handler128    classify_img_button.click(fn=process_img, inputs=upload_img, outputs=output_img)129 130if __name__ == "__main__":131    demo.launch()