File size: 1,729 Bytes
deaa6fa 54e36a3 deaa6fa 42a3cba 54e36a3 deaa6fa |
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 |
import cv2
from ultralytics import YOLO
import gradio as gr
import tempfile
import os
# Load trained model
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()
|