Spaces:
Sleeping
Sleeping
File size: 2,215 Bytes
52b90d0 23eb270 745db4e 23eb270 2d9fc09 23eb270 745db4e 23eb270 7a7dcae 52b90d0 23eb270 |
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 |
import gradio as gr
import numpy as np
import tensorflow as tf
import cv2
import os
import requests
from yolo import Yolo
model_link = "https://intranet-projects-files.s3.amazonaws.com/holbertonschool-ml/yolo.h5"
im_link = "https://miro.medium.com/v2/resize:fit:720/format:webp/1*EYFejGUjvjPcc4PZTwoufw.jpeg"
def download_model():
if not os.path.exists('data'):
os.makedirs('data')
if not os.path.exists('data/yolo.h5'):
print("Downloading model...")
r = requests.get(model_link)
with open('data/yolo.h5', 'wb') as f:
f.write(r.content)
print("Model downnloaded")
else:
print("Model already exists locally.")
if not os.path.exists('data/exm.jpg'):
print("Downloading image...")
r = requests.get(im_link)
with open('data/exm.jpg', 'wb') as f:
f.write(r.content)
print("Image downnloaded")
def run(img):
np.random.seed(0)
anchors = np.array([[[116, 90], [156, 198], [373, 326]],
[[30, 61], [62, 45], [59, 119]],
[[10, 13], [16, 30], [33, 23]]])
yolo = Yolo('data/yolo.h5', 'coco_classes.txt', 0.6, 0.5, anchors)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
boxes, box_classes, box_scores = yolo.predict_frame(img)
for idx, box in enumerate(boxes):
top_left_x = int(box[0])
top_left_y = int(box[1])
bottom_right_x = int(box[2])
bottom_right_y = int(box[3])
class_name = yolo.class_names[box_classes[idx]]
score = box_scores[idx]
color = (255, 0, 0)
cv2.rectangle(img, (top_left_x, top_left_y),
(bottom_right_x, bottom_right_y),
color, 2)
text = f"{class_name} {score:.2f}"
cv2.putText(img, text, (top_left_x, top_left_y - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1,
cv2.LINE_AA)
# Convert image back from RGB to BGR for displaying with OpenCV
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
return img
demo = gr.Interface(run, "image", "image" )
if __name__ == "__main__":
download_model()
demo.launch()
|