File size: 1,620 Bytes
27b0fa8 bcebd9d 27b0fa8 |
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 |
from clearml import Model
from ultralytics import YOLO
from PIL import Image, ImageDraw
from huggingface_hub import hf_hub_download
import gradio as gr
# fetching the model from the ClearML
def fetch_model_from_HG():
try:
hg_model = hf_hub_download(repo_id="manvi23/GroupG", filename="models/best.pt")
return hg_model
except Exception as e:
print(f"Failed to fetch model from HuggingFace: {e}")
raise
model_path = fetch_model_from_HG()
model = YOLO(model_path)
def predict_and_visualize(image, confidence_threshold=0.3):
results = model.predict(image, conf=confidence_threshold)
boxes = results[0].boxes.xyxy.cpu().numpy() # Bounding box coordinates
scores = results[0].boxes.conf.cpu().numpy() # Confidence scores
# Draw bounding boxes on the image
draw = ImageDraw.Draw(image)
for box, score in zip(boxes, scores):
x_min, y_min, x_max, y_max = box
draw.rectangle([x_min, y_min, x_max, y_max], outline="red", width=3)
draw.text((x_min, y_min), f"{score:.2f}", fill="red")
return image
# Gradio interface function
def gradio_app(image):
result_image = predict_and_visualize(image)
return result_image
# Create Gradio app
interface = gr.Interface(
fn=gradio_app,
inputs=[
gr.Image(type="pil"), # Input image
],
outputs=gr.Image(type="pil"), # Output image
title="Object Detection with model of Group G",
description="Upload an image to detect the objects.",
)
# Launch the Gradio app
interface.launch(share=True)
|