import gradio as gr
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import torch

model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

LABELS = [
    "auto-rickshaw on road",
    "cattle crossing road",
    "pothole on road",
    "speed breaker",
    "pedestrian crossing",
    "traffic signal",
    "wrong side driving",
    "narrow lane",
    "highway road",
    "road divider"
]

def detect(image):
    inputs = processor(
        text=LABELS,
        images=image,
        return_tensors="pt",
        padding=True
    )
    outputs = model(**inputs)
    probs = outputs.logits_per_image.softmax(dim=1)
    return {LABELS[i]: float(probs[0][i]) for i in range(len(LABELS))}

demo = gr.Interface(
    fn=detect,
    inputs=gr.Image(type="pil", label="Road Image Upload karo"),
    outputs=gr.Label(num_top_classes=5, label="Detection Results"),
    title="🚗 SamyamLM — Self Driving Car Detector",
    description="Indian road conditions detect karo!"
)

demo.launch()