ObjectDetection / app.py
RakanAlsheraiwi's picture
Update app.py
7d6403f verified
raw
history blame
3.49 kB
import cv2
import torch
from PIL import Image, ImageDraw
import gradio as gr
import pandas as pd
from transformers import pipeline
# تحميل النموذج
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
translator = pipeline("translation_en_to_ar", model="Helsinki-NLP/opus-mt-en-ar")
# دالة لاكتشاف الكائنات في الصور
def detect_and_draw_image(input_image):
results = model(input_image)
detections = results.xyxy[0].numpy()
draw = ImageDraw.Draw(input_image)
counts = {}
for detection in detections:
xmin, ymin, xmax, ymax, conf, class_id = detection
label = model.names[int(class_id)]
counts[label] = counts.get(label, 0) + 1
draw.rectangle([(xmin, ymin), (xmax, ymax)], outline="red", width=2)
draw.text((xmin, ymin), f"{label}: {conf:.2f}", fill="white")
translated_labels = translator(list(counts.keys()))
df = pd.DataFrame({
'Label (English)': list(counts.keys()),
'Label (Arabic)': [t['translation_text'] for t in translated_labels],
'Object Count': list(counts.values())
})
return input_image, df
# دالة لاكتشاف الكائنات في الفيديو
def detect_and_draw_video(video_path):
cap = cv2.VideoCapture(video_path)
frames = []
overall_counts = {}
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
frame = cv2.resize(frame, (640, 480))
results = model(frame)
detections = results.xyxy[0].numpy()
for detection in detections:
xmin, ymin, xmax, ymax, conf, class_id = detection
label = model.names[int(class_id)]
overall_counts[label] = overall_counts.get(label, 0) + 1
cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (255, 0, 0), 2)
cv2.putText(frame, f"{label}: {conf:.2f}", (int(xmin), int(ymin) - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
frames.append(frame)
cap.release()
output_path = 'output.mp4'
out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 20.0, (640, 480))
for frame in frames:
out.write(frame)
out.release()
translated_labels = translator(list(overall_counts.keys()))
df = pd.DataFrame({
'Label (English)': list(overall_counts.keys()),
'Label (Arabic)': [t['translation_text'] for t in translated_labels],
'Object Count': list(overall_counts.values())
})
return output_path, df
# واجهة صورة
image_interface = gr.Interface(
fn=detect_and_draw_image,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=[gr.Image(type="pil"), gr.Dataframe(label="Object Counts")],
title="Object Detection for Images",
description="Upload an image to see the objects detected and their counts.",
examples=['assets/MessiVsAlhilal.jpg', 'assets/Manhattan002_0.webp'] # إضافة الأمثلة هنا
)
video_interface = gr.Interface(
fn=detect_and_draw_video,
inputs=gr.Video(label="Upload Video"),
outputs=[gr.Video(label="Processed Video"), gr.Dataframe(label="Object Counts")],
title="Object Detection for Videos",
description="Upload a video to see the objects detected and their counts."
)
app = gr.TabbedInterface([image_interface, video_interface], ["Image Detection", "Video Detection"])
app.launch(debug=True)