Added images and model with Git LFS

app.py

@@ -34,8 +34,12 @@ def load_model(model_filename):
 def classify_new_image(input_image_path, model_choice):
     """
     Expects input_image_path as a file path and model_choice as one of the keys in models.
-    Loads the image, processes it by extracting patches and computing predictions on each patch,
-    aggregates patch predictions, and returns the final predicted class and probabilities.
+    Loads the image, processes it by extracting patches, classifies each patch, and aggregates
+    patch predictions. Also, draws transparent overlays on each patch according to its predicted label.
+    Returns:
+        annotated_image_rgb (numpy array): The image with transparent overlays.
+        final_prediction (str): The final predicted class.
+        prob_dict (dict): Dictionary of class probabilities.
     """
     global models, training_log, class_names
     progress_log = training_log + "\nStarting classification...\n"
@@ -69,13 +73,12 @@ def classify_new_image(input_image_path, model_choice):
     
     # Patch extraction parameters
     patch_size = (100, 100)
-    edge_density_thresh_low = 0.0
-    edge_density_thresh_high = 0.5
     patch_w, patch_h = patch_size
     img_h, img_w = gray.shape
     
     valid_patch_count = 0
     summed_probabilities = None
+    overlays_list = []  # To store (x, y, w, h, predicted_label) for each valid patch
     
     # Loop over non-overlapping patches
     for y in range(0, img_h - patch_h + 1, patch_h):
@@ -87,12 +90,16 @@ def classify_new_image(input_image_path, model_choice):
             density = num_edge_pixels / total_pixels
             progress_log += f"Patch at ({x}, {y}) - edge density: {density:.3f}\n"
             
-            if edge_density_thresh_low < density < edge_density_thresh_high:
+            if 0.0 < density < 0.5:
                 valid_patch_count += 1
                 features = extract_features_from_image(patch)
                 feature_vector = features['combined_features'].reshape(1, -1)
                 patch_probabilities = classifier.predict_proba(feature_vector)[0]
-                progress_log += f"Patch at ({x}, {y}) predicted probabilities: {patch_probabilities}\n"
+                predicted_index = np.argmax(patch_probabilities)
+                predicted_label = class_names[predicted_index]
+                progress_log += f"Patch at ({x}, {y}) predicted: {predicted_label} with probabilities {patch_probabilities}\n"
+                
+                overlays_list.append((x, y, patch_w, patch_h, predicted_label))
                 
                 if summed_probabilities is None:
                     summed_probabilities = patch_probabilities
@@ -107,10 +114,8 @@ def classify_new_image(input_image_path, model_choice):
         summed_probabilities = classifier.predict_proba(feature_vector)[0]
         valid_patch_count = 1
     
-    # Average the probabilities from all valid patches
+    # Average the probabilities from all valid patches and normalize them
     averaged_probabilities = summed_probabilities / valid_patch_count
-
-    # Normalize the averaged probabilities so they sum to 1
     normalized_probabilities = averaged_probabilities / np.sum(averaged_probabilities)
     
     final_prediction_index = np.argmax(normalized_probabilities)
@@ -120,7 +125,32 @@ def classify_new_image(input_image_path, model_choice):
     
     print(progress_log)
     print(prob_dict)
-    return final_prediction, prob_dict
+    
+    # Create an annotated image with transparent overlays
+    annotated_image = resized_image.copy()
+    overlay = annotated_image.copy()
+    alpha = 0.4  # Transparency factor
+    
+    # Define overlay colors in BGR for each class (adjust as desired)
+    color_map = {
+        'wood': (0, 255, 255),   # Yellow (BGR format)
+        'brick': (0, 0, 255),    # Red (BGR format)
+        'stone': (128, 128, 128) # Gray (BGR format)
+}
+
+    for (x, y, w, h, label) in overlays_list:
+        label_lower = label.lower()
+        if prob_dict[label] < 0.2:
+            continue
+        color = color_map.get(label_lower, (0, 255, 0))  # Default to green if unknown
+        cv2.rectangle(overlay, (x, y), (x+w, y+h), color, thickness=-1)
+    
+    # Blend the overlay with the original image
+    annotated_image = cv2.addWeighted(overlay, alpha, annotated_image, 1 - alpha, 0)
+    # Convert annotated image from BGR to RGB for Gradio display
+    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
+    
+    return final_prediction, prob_dict, annotated_image_rgb
 
 # ---------------------------------------------------------------------
 # Gradio Interface Setup using file paths and model selection
@@ -130,7 +160,6 @@ if __name__ == "__main__":
     load_model(model_filename)
     
     # Create a dropdown for model selection.
-    # If models dict is not yet populated, we set default choices.
     model_choices = list(models.keys()) if models else ['svm', 'rf', 'combined']
     
     iface = gr.Interface(
@@ -141,11 +170,12 @@ if __name__ == "__main__":
         ],
         outputs=[
             gr.Label(label="Predicted Class"),
-            gr.Label(label="Probabilities")
+            gr.Label(label="Probabilities"),
+            gr.Image(label="Annotated Image")
         ],
         title="Stone, Wood, Brick Classifier",
         description=("Upload an image and select a classifier model (svm, rf, combined) to classify it.\n\n"
                      "The image is processed by subdividing it into patches and aggregating the predictions. "
-                     "Progress logs are printed to the terminal.")
+                     "Transparent overlays are drawn on detected objects. Progress logs are printed to the terminal.")
     )
     iface.launch(share=True)