Spaces:
Sleeping
Sleeping
| 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 or KNN Model.""") | |
| 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() | |