Spaces:
Sleeping
Sleeping
| 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" | |
| 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 downloaded") | |
| else: | |
| print("Model already exists locally.") | |
| 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() | |