File size: 2,560 Bytes
edf0436 |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
from ultralytics import YOLO
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
import os
import tempfile
import cv2
model_path = "best.onnx"
def load_model(repo_id):
download_dir = snapshot_download(repo_id)
print(download_dir)
path = os.path.join(download_dir, "best.onnx")
print(path)
detection_model = YOLO(path, task='detect')
return detection_model
def process_image(pilimg):
source = pilimg
result = detection_model.predict(source, conf=0.5)
img_bgr = result[0].plot()
out_pilimg = Image.fromarray(img_bgr[..., ::-1]) # RGB-order PIL image
return out_pilimg
def process_video(video):
cap = cv2.VideoCapture(video)
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
fps = cap.get(cv2.CAP_PROP_FPS)
temp_dir = tempfile.mkdtemp()
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
output_path = os.path.join(temp_dir, "output.mp4")
output = cv2.VideoWriter(output_path, fourcc, fps, (int(width), int(height)))
# Loop through the video frames
while cap.isOpened():
# Read a frame from the video
success, frame = cap.read()
if success:
# Run YOLO inference on the frame on GPU Device 0
results = detection_model.predict(frame, conf=0.5)
# Visualize the results on the frame
annotated_frame = results[0].plot()
# Write the annotated frame
output.write(annotated_frame)
output.release()
output.release()
cv2.destroyAllWindows()
cv2.waitKey(1)
return output_path
REPO_ID = "1657866Y/grocery"
detection_model = load_model(REPO_ID)
# Create the interface for image upload
image_interface = gr.Interface(fn=process_image,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"))
# Create the interface for video upload
video_interface = gr.Interface(fn=process_video,
inputs=gr.Video(label="Upload a Video"),
outputs="video")
# Use gr.Blocks to arrange components and launch the app
with gr.Blocks() as app:
# Add a header using Markdown
gr.Markdown("# Grocery? No wait!")
gr.Markdown("Choose whether to upload an image or a video below!")
# Add the tabbed interface
gr.TabbedInterface([image_interface, video_interface],
tab_names=["Image Upload", "Video Upload"])
# Launch the interface
app.launch() |