File size: 2,924 Bytes
8f970ae
335d300
43a659b
534f8cc
43a659b
 
335d300
ff8d4cb
43a659b
335d300
 
43a659b
335d300
 
534f8cc
335d300
4e01a45
534f8cc
335d300
 
534f8cc
335d300
 
 
534f8cc
335d300
43a659b
335d300
 
43a659b
335d300
 
52ccd52
335d300
4e01a45
335d300
 
4e01a45
 
335d300
43a659b
 
8f970ae
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
335d300
 
8f970ae
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
88
89
90
91
"""import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np

# Load the YOLO model
MODEL_URL = "https://huggingface.co/ayoubsa/yolo_model/resolve/main/best.pt"
model = YOLO(MODEL_URL)

# Define the prediction function
def predict(input_img):
    try:
        # Convert PIL Image to NumPy array
        image_array = np.array(input_img)

        # Perform inference
        results = model(image_array)

        # Extract detected class names
        detected_classes = [model.names[int(cls)] for cls in results[0].boxes.cls]

        # Render results on the image
        results[0].plot()  # Render bounding boxes on the image
        output_image = Image.fromarray(results[0].orig_img)

        return output_image, {cls: 1.0 for cls in detected_classes}  # Dummy scores for simplicity
    except Exception as e:
        print(f"Error during processing: {e}")
        return None, {"Error": str(e)}

# Gradio app configuration
gradio_app = gr.Interface(
    predict,
    inputs=gr.Image(label="Upload an Image", type="pil"),
    outputs=[
        gr.Image(label="Predicted Image with Bounding Boxes"),  # Rendered image with bounding boxes
        gr.Label(label="Detected Classes"),  # Detected class names
    ],
    title="YOLO Object Detection App",
    description="Upload an image, and the YOLO model will detect objects in it.",
)

if __name__ == "__main__":
    gradio_app.launch()"""

import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np

# Load the YOLO model
MODEL_URL = "https://huggingface.co/ayoubsa/yolo_model/resolve/main/best.pt"
model = YOLO(MODEL_URL)

# Define the prediction function
def predict(input_img):
    try:
        # Convert PIL Image to NumPy array
        image_array = np.array(input_img)

        # Perform inference
        results = model(image_array)

        # Extract detected class names
        detected_classes = [model.names[int(cls)] for cls in results[0].boxes.cls]

        # Render results on the image
        rendered_image = results[0].plot()  # This method returns the image with bounding boxes
        output_image = Image.fromarray(rendered_image)  # Convert the rendered image to a PIL Image

        return output_image, {cls: 1.0 for cls in detected_classes}  # Dummy scores for simplicity
    except Exception as e:
        print(f"Error during processing: {e}")
        return None, {"Error": str(e)}

# Gradio app configuration
gradio_app = gr.Interface(
    predict,
    inputs=gr.Image(label="Upload an Image", type="pil"),
    outputs=[
        gr.Image(label="Predicted Image with Bounding Boxes"),  # Rendered image with bounding boxes
        gr.Label(label="Detected Classes"),  # Detected class names
    ],
    title="YOLO Object Detection App",
    description="Upload an image, and the YOLO model will detect objects in it.",
)

if __name__ == "__main__":
    gradio_app.launch()