Spaces:
Running
Running
File size: 2,164 Bytes
8eb3cb3 77c9c0f 8eb3cb3 77c9c0f 8eb3cb3 77c9c0f dd65482 8eb3cb3 dd65482 77c9c0f dd65482 77c9c0f 8eb3cb3 77c9c0f dd65482 77c9c0f |
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 |
import os
import gradio as gr
import numpy as np
from PIL import Image
from ultralytics import YOLO
# Load the model
model = YOLO('best.pt')
# Path to the photos folder
photos_folder = "Photos"
def load_images_from_folder(folder):
images = []
for filename in os.listdir(folder):
if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
img_path = os.path.join(folder, filename)
img = Image.open(img_path)
images.append((img, filename))
return images
def predict(image):
try:
image = np.array(image)
results = model(image)
result_image = results[0].plot()
return Image.fromarray(result_image)
except Exception as e:
print(f"Error during prediction: {e}")
return "Error"
def load_image_from_gallery(images, index):
if images and 0 <= index < len(images):
image = images[index]
if isinstance(image, tuple):
image = image[0]
return image
return None
def gallery_click_event(images, evt: gr.SelectData):
index = evt.index
selected_img = load_image_from_gallery(images, index)
return selected_img
def clear_image():
return None
# Load images at the start
images = load_images_from_folder(photos_folder)
with gr.Blocks(css=".container { background-color: white; }") as demo:
with gr.Row():
with gr.Column():
selected_image = gr.Image(label="Selected Image from Gallery", type="pil")
clear_button = gr.Button("Clear")
with gr.Column():
image_gallery = gr.Gallery(label="Image Gallery", elem_id="gallery", type="pil", value=[img for img, _ in images])
with gr.Column():
result_image = gr.Image(label="Result Image", type="pil")
image_gallery.select(
fn=gallery_click_event,
inputs=image_gallery,
outputs=selected_image
)
selected_image.change(
fn=predict,
inputs=selected_image,
outputs=result_image
)
clear_button.click(
fn=clear_image,
inputs=None,
outputs=selected_image
)
demo.launch()
|