Spaces:
Sleeping
Sleeping
File size: 3,271 Bytes
e76a638 |
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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import gradio as gr
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from model import BoundingBoxPredictor
import matplotlib.patches as patches
import os
import uuid
predictor = BoundingBoxPredictor()
current_dir = os.path.dirname(os.path.abspath(__file__))
cnn_path = os.path.join(current_dir, 'convolutional_nn.h5')
knn_path = os.path.join(current_dir, 'knn_model_tuned.pkl')
sample_images = [
os.path.join(current_dir, f"image{i}.jpg")
for i in range(4996, 5000)
]
try:
predictor.load_models(cnn_path, knn_path)
print("Models loaded successfully!")
except Exception as e:
print(f"Error loading models: {str(e)}")
raise
def predict_and_draw_bbox(input_image, model_choice):
if input_image is None:
return None, "Please upload an image."
try:
bbox = predictor.predict(input_image, model_type=model_choice.lower())
img_array = np.array(input_image)
fig, ax = plt.subplots()
ax.imshow(img_array, cmap='gray')
rect = patches.Rectangle(
(bbox['x'], bbox['y']),
bbox['width'],
bbox['height'],
linewidth=2,
edgecolor='r',
facecolor='none'
)
ax.add_patch(rect)
plt.axis('off')
output_filename = f'output_{uuid.uuid4().hex[:8]}.png'
output_path = os.path.join(current_dir, output_filename)
plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
plt.close()
return (
output_path,
f"Model Used: {model_choice}\n\nBounding Box Coordinates:\nX: {bbox['x']:.2f}\nY: {bbox['y']:.2f}\nWidth: {bbox['width']:.2f}\nHeight: {bbox['height']:.2f}"
)
except Exception as e:
return None, f"Error processing image: {str(e)}"
def select_sample(evt: gr.SelectData):
selected_image = Image.open(sample_images[evt.index])
return selected_image
with gr.Blocks() as iface:
gr.Markdown("# Bounding Box Detector")
gr.Markdown("""Upload an image to detect the bounding box. Choose between:
- CNN Model (Best performance, IoU: 0.65)
- KNN Model (Faster, IoU: 0.47)
The models work best with grayscale images of size 128x128.""")
with gr.Row():
with gr.Column():
input_image = gr.Image(type="pil", label="Input Image")
model_choice = gr.Radio(
choices=["CNN", "KNN"],
label="Choose Model",
value="CNN"
)
submit_btn = gr.Button("Detect Bounding Box")
with gr.Column():
output_image = gr.Image(type="filepath", label="Detected Bounding Box")
output_text = gr.Textbox(label="Coordinates")
gr.Markdown("### Sample Images (click to use)")
gallery = gr.Gallery(
value=sample_images,
label="Sample Images",
show_label=False,
columns=4,
height="auto"
).select(select_sample, None, input_image)
submit_btn.click(
predict_and_draw_bbox,
inputs=[input_image, model_choice],
outputs=[output_image, output_text]
)
if __name__ == "__main__":
iface.launch()
|