Spaces:
Sleeping
Sleeping
File size: 3,384 Bytes
af0a0f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import os
import tempfile
import cv2
import numpy as np
import gradio as gr
import supervision as sv
from ultralytics import YOLO
# 1. Cấu hình Model
MODEL_PATH = "best.pt" # Thay bằng model đã huấn luyện
model = YOLO(MODEL_PATH)
CLASS_NAMES_DICT = model.names
SELECTED_CLASS_NAMES = list(CLASS_NAMES_DICT.values())
# Khởi tạo Annotators
box_annotator = sv.BoxAnnotator(thickness=2)
label_annotator = sv.LabelAnnotator(text_thickness=1, text_scale=0.8)
def process_image(image):
"""Xử lý ảnh đơn lẻ."""
if image is None:
return None
results = model(image, conf=0.25, verbose=False)[0]
detections = sv.Detections.from_ultralytics(results)
labels = [
f"{CLASS_NAMES_DICT[class_id]} {conf:0.2f}"
for conf, class_id in zip(detections.confidence, detections.class_id)
]
annotated_image = image.copy()
annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections)
annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
return annotated_image
def process_video(video_path):
"""Xử lý video."""
if video_path is None: return None
output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
def callback(frame: np.ndarray, index: int) -> np.ndarray:
results = model(frame, conf=0.25, verbose=False)[0]
detections = sv.Detections.from_ultralytics(results)
labels = [
f"{CLASS_NAMES_DICT[class_id]} {conf:0.2f}"
for conf, class_id in zip(detections.confidence, detections.class_id)
]
annotated_frame = frame.copy()
annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections)
annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
return annotated_frame
sv.process_video(source_path=video_path, target_path=output_path, callback=callback)
return output_path
# --- Giao diện Gradio ---
with gr.Blocks(title="Road Damage", theme=gr.themes.Soft()) as demo:
gr.Markdown("# 🚧 Road Damage Detection")
gr.Markdown("Hệ thống nhận diện và đánh giá xuống cấp mặt đường.")
with gr.Tabs():
# Tab Xử lý Ảnh
with gr.TabItem("🖼️ Xử lý Ảnh"):
with gr.Row():
img_input = gr.Image(label="Tải ảnh lên")
img_output = gr.Image(label="Kết quả phân tích")
img_btn = gr.Button("Phân tích Ảnh", variant="primary")
img_btn.click(fn=process_image, inputs=img_input, outputs=img_output)
# Tab Xử lý Video
with gr.TabItem("🎥 Xử lý Video"):
with gr.Row():
vid_input = gr.Video(label="Tải video lên")
vid_output = gr.Video(label="Video kết quả")
vid_btn = gr.Button("Phân tích Video", variant="primary")
vid_btn.click(fn=process_video, inputs=vid_input, outputs=vid_output)
gr.Markdown("---")
gr.Markdown("""
### Thành viên:
- An Hoàng Anh - 223332813
- Trần Hải Nam - 223332840
- Đoàn Minh Thành - 223332848
- Nguyễn Công Thành - 223332850
""")
if __name__ == "__main__":
demo.launch() |