Pooja-S-Hub/Components_detection

1import gradio as gr2from ultralytics import YOLO3import cv24import json5import numpy as np6import pandas as pd7from datetime import datetime8import os9 10# Load YOLO model11model = YOLO("robot_yolo_medium_v2.pt")12 13# Load all components14with open("components.json") as f:15    COMPONENTS = json.load(f)["components"]16 17# Create folder to save Excel files18if not os.path.exists("reports"):19    os.makedirs("reports")20 21def detect_missing(image):22    results = model.predict(image, imgsz=640)23    detected_classes = []24 25    # Collect detected component names26    for r in results:27        for box in r.boxes:28            class_id = int(box.cls[0])29            detected_classes.append(COMPONENTS[class_id])30 31    # Identify missing components32    missing = [c for c in COMPONENTS if c not in detected_classes]33 34    # Draw bounding boxes on image35    img = image.copy()36    for r in results:37        for box in r.boxes:38            x1, y1, x2, y2 = map(int, box.xyxy[0])39            cls_id = int(box.cls[0])40            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)41            cv2.putText(img, COMPONENTS[cls_id], (x1, y1-10),42                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)43 44    # Save missing info to Excel45    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")46    df = pd.DataFrame({"Missing Components": missing})47    excel_path = f"reports/missing_{timestamp}.xlsx"48    df.to_excel(excel_path, index=False)49 50    return img, ", ".join(missing) if missing else "All components detected!"51 52# Gradio interface53iface = gr.Interface(54    fn=detect_missing,55    inputs=gr.Image(type="numpy"),56    outputs=[gr.Image(type="numpy"), gr.Textbox()],57    title="Robot Component Checker",58    description="Upload a photo of robot components. Missing components will be flagged and saved."59)60 61iface.launch()62