|
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") |
|
|
|
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() |
|
|
