Upload 7 files

app.py

@@ -0,0 +1,107 @@
+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()

model.py

@@ -0,0 +1,108 @@
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import joblib
+import os
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, Flatten
+
+class BoundingBoxPredictor:
+    def __init__(self):
+        self.cnn_model = None
+        self.knn_model = None
+        self.input_shape = (128, 128)
+        # Updated normalization values based on actual dataset analysis
+        self.y_min = np.array([0., 0., 30., 30.])  # min values for x, y, width, height
+        self.y_max = np.array([97., 97., 49., 49.])  # max values for x, y, width, height
+        
+    def create_cnn_model(self):
+        model = Sequential([
+            Conv2D(32, (3, 3), activation='relu', input_shape=(*self.input_shape, 1)),
+            MaxPooling2D((2, 2)),
+            Conv2D(64, (3, 3), activation='relu'),
+            MaxPooling2D((2, 2)),
+            Flatten(),
+            Dense(128, activation='relu'),
+            Dense(4, activation='sigmoid')
+        ])
+        model.compile(optimizer='adam', loss='mse', metrics=['mse'])
+        return model
+        
+    def load_models(self, cnn_path, knn_path):
+        # Check if files exist
+        if not os.path.exists(cnn_path):
+            raise FileNotFoundError(f"CNN model file not found: {cnn_path}")
+        if not os.path.exists(knn_path):
+            raise FileNotFoundError(f"KNN model file not found: {knn_path}")
+            
+        print(f"Loading CNN model from: {cnn_path}")
+        try:
+            self.cnn_model = tf.keras.models.load_model(cnn_path)
+        except Exception as e:
+            print(f"Error loading CNN model: {str(e)}")
+            print("Creating new CNN model...")
+            self.cnn_model = self.create_cnn_model()
+            try:
+                self.cnn_model.load_weights(cnn_path)
+                print("Successfully loaded CNN weights")
+            except:
+                print("Could not load CNN weights, using uninitialized model")
+        
+        print("CNN model ready")
+        
+        print(f"Loading KNN model from: {knn_path}")
+        self.knn_model = joblib.load(knn_path)
+        print("KNN model loaded successfully")
+        
+    def preprocess_image(self, image):
+        # Convert to grayscale if needed
+        if image.mode != 'L':
+            image = image.convert('L')
+            
+        # Resize image
+        image = image.resize(self.input_shape)
+        
+        # Convert to numpy array and normalize
+        img_array = np.array(image)
+        img_array = img_array / 255.0
+        
+        return img_array
+        
+    def predict(self, image, model_type='cnn'):
+        # Check if models are loaded
+        if model_type == 'cnn' and self.cnn_model is None:
+            raise ValueError("CNN model not loaded. Please call load_models first.")
+        if model_type == 'knn' and self.knn_model is None:
+            raise ValueError("KNN model not loaded. Please call load_models first.")
+            
+        # Preprocess the image
+        processed_image = self.preprocess_image(image)
+        
+        try:
+            if model_type == 'cnn':
+                # Reshape for CNN input
+                img_array = processed_image.reshape(1, *self.input_shape, 1)
+                # Get normalized predictions (between 0 and 1)
+                prediction = self.cnn_model.predict(img_array)[0]
+                # Denormalize predictions to original scale
+                prediction = prediction * (self.y_max - self.y_min) + self.y_min
+            else:  # KNN
+                # Flatten for KNN input
+                img_array = processed_image.flatten().reshape(1, -1)
+                # Get normalized predictions
+                prediction = self.knn_model.predict(img_array)[0]
+                # Denormalize predictions to original scale
+                prediction = prediction * (self.y_max - self.y_min) + self.y_min
+            
+            # Ensure predictions are within valid ranges
+            prediction = np.clip(prediction, self.y_min, self.y_max)
+            
+            return {
+                'x': float(prediction[0]),
+                'y': float(prediction[1]),
+                'width': float(prediction[2]),
+                'height': float(prediction[3])
+            }
+        except Exception as e:
+            print(f"Error during prediction with {model_type.upper()} model: {str(e)}")
+            raise

requirements.txt

@@ -0,0 +1,8 @@
+tensorflow>=2.0.0
+numpy
+Pillow
+gradio>=4.0.0
+matplotlib
+scikit-learn
+joblib
+huggingface-hub