File size: 2,269 Bytes
54e76cb |
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 |
import streamlit as st
import cv2
import tempfile
import os
import numpy as np
from ultralytics import YOLO
from PIL import Image
def process_video(video_path, model):
cap = cv2.VideoCapture(video_path)
temp_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(temp_output.name, fourcc, cap.get(cv2.CAP_PROP_FPS),
(int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
results = model(frame)
for result in results:
for box in result.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
label = result.names[int(box.cls[0])]
conf = float(box.conf[0])
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, f'{label}: {conf:.2f}', (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
out.write(frame)
cap.release()
out.release()
return temp_output.name
def main():
st.set_page_config(page_title="Bottle Label Checker", page_icon="🍾")
st.title("Bottle Label Checking System using YOLO & Gemini")
uploaded_video = st.file_uploader("Upload a video", type=["mp4", "avi", "mov", "mkv"])
if uploaded_video is not None:
temp_video_path = os.path.join(tempfile.gettempdir(), uploaded_video.name)
with open(temp_video_path, "wb") as f:
f.write(uploaded_video.read())
st.video(temp_video_path)
model = YOLO("yolov8n.pt") # Load YOLO model
if st.button("Process Video"):
st.write("Processing video... This may take some time.")
output_path = process_video(temp_video_path, model)
st.video(output_path)
st.success("Processing complete!")
with open(output_path, "rb") as file:
st.download_button("Download Processed Video", file, file_name="processed_video.mp4", mime="video/mp4")
if __name__ == "__main__":
main()
|