|
|
import cv2 |
|
|
from ultralytics import YOLO |
|
|
import gradio as gr |
|
|
import tempfile |
|
|
import os |
|
|
|
|
|
|
|
|
model = YOLO("best.pt") |
|
|
|
|
|
def process_video(input_video): |
|
|
cap = cv2.VideoCapture(input_video) |
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v') |
|
|
fps = int(cap.get(cv2.CAP_PROP_FPS)) |
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|
|
|
|
|
tmp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") |
|
|
output_path = tmp_out.name |
|
|
tmp_out.close() |
|
|
|
|
|
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) |
|
|
|
|
|
while True: |
|
|
ret, frame = cap.read() |
|
|
if not ret: |
|
|
break |
|
|
|
|
|
results = model(frame)[0] |
|
|
for box in results.boxes: |
|
|
x1, y1, x2, y2 = map(int, box.xyxy[0]) |
|
|
conf = box.conf[0] |
|
|
if conf < 0.4: |
|
|
continue |
|
|
cls = int(box.cls[0]) |
|
|
label = f"Cleaner {conf:.2f}" |
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 200, 0), 2) |
|
|
cv2.putText(frame, label, (x1, y1 - 10), |
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 200, 0), 2) |
|
|
|
|
|
out.write(frame) |
|
|
|
|
|
cap.release() |
|
|
out.release() |
|
|
cv2.destroyAllWindows() |
|
|
|
|
|
return output_path |
|
|
|
|
|
with gr.Blocks(css=".gradio-container {color: #cce0ff;}") as demo: |
|
|
gr.Markdown("## Detecting hands with cleaning cloth") |
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
inp = gr.Video(label="Upload Video") |
|
|
btn = gr.Button("Process Video", variant="primary") |
|
|
with gr.Column(): |
|
|
out = gr.Video(label="Output Video") |
|
|
btn.click(fn=process_video, inputs=inp, outputs=out) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch() |
|
|
|
