File size: 1,270 Bytes
230fbb8
78117b4
 
230fbb8
78117b4
230fbb8
10183ac
230fbb8
 
 
 
 
10183ac
 
 
230fbb8
 
 
 
 
 
 
 
b489f83
230fbb8
 
 
 
9313dfd
230fbb8
 
 
78117b4
230fbb8
 
 
10183ac
230fbb8
 
 
78117b4
 
230fbb8
 
 
 
 
 
 
 
 
 
 
 
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
import gradio as gr
from PIL import Image
import numpy as np
from ultralytics import YOLO
import os

def handle_classify(image=None):
    """This function performs YOLOv8 object detection on the given image.

    Args:
        image (gr.inputs.Image, optional): Input image to detect objects on. Defaults to None.
    """
    
    if not image:
        return "No image found"

    model_path = "racist2.0.pt"
    model = YOLO(model_path)

    results = model(image)
    
    result = results[0]
    
    top5 = [[result.names[class_index], str(round(result.probs.top5conf.tolist()[rank], 4)*100)+'%']
                    for class_index, rank in zip(result.probs.top5, range(5))]
    
    print(top5)
    
    return "\n".join(["\t".join(row) for row in top5])


inputs = [
    gr.Image(type='numpy', label="Input Image"),
]


outputs = gr.Textbox()

title = "Racist model v2"

SAMPLE_DIR = 'samples'
examples = [np.array(Image.open(os.path.join(SAMPLE_DIR, path))) for path in os.listdir(SAMPLE_DIR)]

yolo_app = gr.Interface(
    fn=handle_classify,
    inputs=inputs,
    outputs=outputs,
    title=title,
    examples=examples,
    cache_examples=True,
)

# Launch the Gradio interface in debug mode with queue enabled
yolo_app.launch(debug=True, enable_queue=True)