LicenseDetector / app.py
fadzwan's picture
Update app.py
638de5c verified
import gradio as gr
from PIL import Image
import numpy as np
import cv2
from ultralytics import YOLO
from transformers import pipeline
import matplotlib.pyplot as plt
# Load models
yolo_model = YOLO("yolov8n.pt")
damage_pipe = pipeline("image-classification", model="beingamit99/car_damage_detection")
# License Plate Detection
def detect_license_plate(img: Image.Image):
results = yolo_model.predict(img)
result = results[0]
img_array = np.array(img.convert("RGB"))
for box in result.boxes.xyxy:
x1, y1, x2, y2 = map(int, box)
img_array = cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
return Image.fromarray(img_array)
# Car Damage Classification
def classify_damage(image: Image.Image):
results = damage_pipe(image)
labels = [res["label"] for res in results]
scores = [res["score"] for res in results]
fig, ax = plt.subplots()
ax.barh(labels, scores, color="crimson")
ax.set_xlabel("Confidence")
ax.set_xlim(0, 1)
ax.set_title("Damage Type Classification")
plt.tight_layout()
return fig
# Build UI
with gr.Blocks() as app:
gr.Markdown("# 🚘 Car Analyzer\nChoose a tool below:")
with gr.Tab("πŸ›‚ License Plate Detection"):
with gr.Row():
with gr.Column():
lp_input = gr.Image(type="pil", label="Upload Car Image")
lp_btn = gr.Button("Detect")
with gr.Column():
lp_output = gr.Image(label="Detected License Plate")
lp_btn.click(detect_license_plate, inputs=lp_input, outputs=lp_output)
with gr.Tab("πŸ’₯ Car Damage Classification"):
with gr.Row():
with gr.Column():
dmg_input = gr.Image(type="pil", label="Upload Damaged Car Image")
dmg_btn = gr.Button("Classify")
with gr.Column():
dmg_output = gr.Plot(label="Classification Results")
dmg_btn.click(classify_damage, inputs=dmg_input, outputs=dmg_output)
app.launch()