done

.gitignore

@@ -0,0 +1,2 @@
+.env
+.DS_Store

app.py

@@ -0,0 +1,60 @@
+import gradio as gr
+import torch
+from PIL import Image, ImageDraw
+from transformers import AutoImageProcessor, ViTForImageClassification
+from ultralytics import YOLO
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+HF_TOKEN = os.getenv("HF_TOKEN") 
+
+processor = AutoImageProcessor.from_pretrained("Abduqayum/Vehicle-Color-Recognition", token=HF_TOKEN)
+model = ViTForImageClassification.from_pretrained("Abduqayum/Vehicle-Color-Recognition", token=HF_TOKEN)
+model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+detection_model = YOLO("yolov8n.pt")  
+
+CLASSES_ID = [2, 3, 5, 7]
+
+def classify_image(image):
+    image = image.convert("RGB")
+
+    results = detection_model(image, classes=CLASSES_ID, conf=0.5, verbose=False)
+    detections = results[0].boxes.data.cpu().numpy()  
+
+    if len(detections) == 0:
+        return "No vehicle detected (car, bus, or truck).", image
+
+    largest = max(detections, key=lambda det: (det[2] - det[0]) * (det[3] - det[1]))
+    x1, y1, x2, y2, conf, cls_id = largest
+    x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
+
+    cropped = image.crop((x1, y1, x2, y2))
+
+    inputs = processor(cropped, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits = outputs.logits
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    pred_idx = probs.argmax(dim=-1).item()
+    label = model.config.id2label[pred_idx]
+    confidence = probs[0, pred_idx].item()
+
+    draw = ImageDraw.Draw(image)
+    draw.rectangle([x1, y1, x2, y2], outline="red", width=4)
+    draw.text((x1, y1 - 10), f"{label} ({confidence:.2%})", fill="red")
+
+    return f"Prediction: {label} (confidence: {confidence:.2%})", image
+
+gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil"),
+    outputs=["text", "image"],
+    title="YOLO Vehicle Detector + ViT Classifier to identify color of vehicles"
+).launch()

requirements.txt

@@ -0,0 +1,5 @@
+ultralytics==8.3.33
+transformers==4.49.0
+pillow==10.4.0
+gradio
+python-dotenv