app.py

import os
import torch
from ultralytics.nn.tasks import DetectionModel
torch.serialization.add_safe_globals([DetectionModel])  # FIX for YOLOv8 load

from ultralytics import YOLO
import gradio as gr
import requests

# Load lightweight YOLOv8 nano model
model = YOLO("yolov8n.pt")

# --- Digi-Key API setup ---
DIGIKEY_CLIENT_ID = os.getenv("DIGIKEY_CLIENT_ID","K9d4a2AaGwQcoAvdNDZVYEOB3sqL4bMg")
DIGIKEY_CLIENT_SECRET = os.getenv("DIGIKEY_CLIENT_SECRET","NxzuxY67eJssGDkA")
DIGIKEY_REDIRECT_URI = os.getenv("DIGIKEY_REDIRECT_URI", "https://nawal20-circuit_ai.hf.space")
DIGIKEY_AUTH_URL = "https://api.digikey.com/v1/oauth2/token"
DIGIKEY_SEARCH_URL = "https://api.digikey.com/Search/v3/Products/Keyword"

# Get token (simplified client credentials flow)
def get_access_token():
    data = {
        "client_id": DIGIKEY_CLIENT_ID,
        "client_secret": DIGIKEY_CLIENT_SECRET,
        "grant_type": "client_credentials",
        "redirect_uri": DIGIKEY_REDIRECT_URI,
    }
    resp = requests.post(DIGIKEY_AUTH_URL, data=data)
    resp.raise_for_status()
    return resp.json().get("access_token")

# Search Digi-Key for component
def search_digikey(query):
    token = get_access_token()
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    payload = {"Keywords": query, "RecordCount": 3}
    resp = requests.post(DIGIKEY_SEARCH_URL, headers=headers, json=payload)
    if resp.status_code != 200:
        return [{"error": f"Failed to fetch for {query}"}]
    data = resp.json()
    results = []
    for p in data.get("Products", []):
        results.append({
            "ManufacturerPartNumber": p.get("ManufacturerPartNumber"),
            "Manufacturer": p.get("Manufacturer", {}).get("Value"),
            "Description": p.get("Description"),
            "Datasheet": p.get("PrimaryDatasheet"),
        })
    return results or [{"error": f"No results for {query}"}]

# Component detection + Digi-Key lookup
def process_circuit(img):
    results = model(img)   # list of Results
    r = results[0]         # take first result
    
    detections = []
    for box in r.boxes:
        cls_id = int(box.cls.cpu().numpy())
        conf = float(box.conf.cpu().numpy())
        if conf >= 0.5:  # filter low confidence detections
            part = model.names[cls_id]
            detections.append({
                "component": part,
                "confidence": round(conf, 2),
                "digikey": search_digikey(part)
            })

    if not detections:
        return {"message": "No components detected"}, img

    return {"components": detections}, r.plot()

demo = gr.Interface(
    fn=process_circuit,
    inputs="image",
    outputs=[gr.JSON(label="Detected Components + Digi-Key Recommendations"),
             gr.Image(label="Annotated Circuit")],
    title="Circuit Component Advisor",
    description="Upload a circuit diagram → Detect components → Suggest Digi-Key parts"
)

demo.launch()