try a version from gemini

app.py

@@ -2,33 +2,47 @@ import gradio as gr
 import numpy as np
 import pandas as pd
 from PIL import Image
+import logging
 
-# Your helper imports and tensorflow models assumed to be loaded here:
+# Your helper imports and tensorflow models are assumed to be in the same directory.
+# Ensure 'clustering.py' and 'utils.py' are present in your HuggingFace Space.
 import clustering
 import utils
 from tensorflow import keras
-import logging
+
+# --- Basic Setup ---
 logging.getLogger().setLevel(logging.INFO)
 
-# Paths to save outputs
+# --- Constants and Model Loading ---
 IMAGE_PATH = "classified_damage_sites.png"
 CSV_PATH = "classified_damage_sites.csv"
 
-# Load models once (adjust filenames as needed)
-model1 = keras.models.load_model('rwthmaterials_dp800_network1_inclusion.h5')
-model2 = keras.models.load_model('rwthmaterials_dp800_network2_damage.h5')
+# Load models once at startup to improve performance
+try:
+    model1 = keras.models.load_model('rwthmaterials_dp800_network1_inclusion.h5')
+    model2 = keras.models.load_model('rwthmaterials_dp800_network2_damage.h5')
+except Exception as e:
+    logging.error(f"Error loading models: {e}")
+    # If models can't load, you might want to stop the app from launching
+    # or display an error message in the UI.
+    raise
 
 damage_classes = {3: "Martensite", 2: "Interface", 0: "Notch", 1: "Shadowing"}
 model1_windowsize = [250, 250]
 model2_windowsize = [100, 100]
 
+
+# --- Core Processing Function (Your original logic) ---
 def damage_classification(SEM_image, image_threshold, model1_threshold, model2_threshold):
+    """
+    This function contains the core scientific logic for classifying damage sites.
+    It returns the classified image and paths to the output files.
+    """
     if SEM_image is None:
-        logging.error("No image provided")
-        return None, None, None
+        # This error will be displayed nicely in the Gradio interface
+        raise gr.Error("Please upload an SEM Image before running classification.")
 
     damage_sites = {}
-
     # Step 1: Clustering to find damage centroids
     all_centroids = clustering.get_centroids(
         SEM_image,
@@ -36,18 +50,17 @@ def damage_classification(SEM_image, image_threshold, model1_threshold, model2_t
         fill_holes=True,
         filter_close_centroids=True,
     )
-
     for c in all_centroids:
         damage_sites[(c[0], c[1])] = "Not Classified"
 
     # Step 2: Model 1 to identify inclusions
-    images_model1 = utils.prepare_classifier_input(SEM_image, all_centroids, window_size=model1_windowsize)
-    y1_pred = model1.predict(np.asarray(images_model1, dtype=float))
-    inclusions = np.where(y1_pred[:, 0] > model1_threshold)[0]
-
-    for idx in inclusions:
-        coord = all_centroids[idx]
-        damage_sites[(coord[0], coord[1])] = "Inclusion"
+    if len(all_centroids) > 0:
+        images_model1 = utils.prepare_classifier_input(SEM_image, all_centroids, window_size=model1_windowsize)
+        y1_pred = model1.predict(np.asarray(images_model1, dtype=float))
+        inclusions = np.where(y1_pred[:, 0] > model1_threshold)[0]
+        for idx in inclusions:
+            coord = all_centroids[idx]
+            damage_sites[(coord[0], coord[1])] = "Inclusion"
 
     # Step 3: Model 2 to classify remaining damage types
     centroids_model2 = [list(k) for k, v in damage_sites.items() if v == "Not Classified"]
@@ -55,7 +68,6 @@ def damage_classification(SEM_image, image_threshold, model1_threshold, model2_t
         images_model2 = utils.prepare_classifier_input(SEM_image, centroids_model2, window_size=model2_windowsize)
         y2_pred = model2.predict(np.asarray(images_model2, dtype=float))
         damage_index = np.asarray(y2_pred > model2_threshold).nonzero()
-
         for i in range(len(damage_index[0])):
             sample_idx = damage_index[0][i]
             class_idx = damage_index[1][i]
@@ -64,6 +76,7 @@ def damage_classification(SEM_image, image_threshold, model1_threshold, model2_t
             damage_sites[(coord[0], coord[1])] = label
 
     # Step 4: Draw boxes on image and save output image
+    # The utils.show_boxes function is assumed to return a PIL Image object
     image_with_boxes = utils.show_boxes(SEM_image, damage_sites, save_image=True, image_path=IMAGE_PATH)
 
     # Step 5: Export CSV file
@@ -74,24 +87,56 @@ def damage_classification(SEM_image, image_threshold, model1_threshold, model2_t
     return image_with_boxes, IMAGE_PATH, CSV_PATH
 
 
+# --- Gradio Interface Definition ---
 with gr.Blocks() as app:
     gr.Markdown("# Damage Classification in Dual Phase Steels")
-
-    image_input = gr.Image(label="Upload SEM Image")
-    cluster_threshold_input = gr.Number(value=20, label="Cluster Threshold")
-    model1_threshold_input = gr.Number(value=0.7, label="Model 1 Threshold")
-    model2_threshold_input = gr.Number(value=0.5, label="Model 2 Threshold")
-
-    output_image = gr.Image(label="Classified Image")
-    download_image_btn = gr.DownloadButton(label="Download Image")
-    download_csv_btn = gr.DownloadButton(label="Download CSV")
-
-    classify_btn = gr.Button("Run Classification")
+    gr.Markdown("Upload a Scanning Electron Microscope (SEM) image and set the thresholds to classify material damage.")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Upload SEM Image")
+            cluster_threshold_input = gr.Number(value=20, label="Image Binarization Threshold")
+            model1_threshold_input = gr.Number(value=0.7, label="Inclusion Model Certainty (0-1)")
+            model2_threshold_input = gr.Number(value=0.5, label="Damage Model Certainty (0-1)")
+            classify_btn = gr.Button("Run Classification", variant="primary")
+
+        with gr.Column(scale=2):
+            output_image = gr.Image(label="Classified Image")
+            # Initialize DownloadButtons as hidden. They will become visible after a successful run.
+            download_image_btn = gr.DownloadButton(label="Download Image", visible=False)
+            download_csv_btn = gr.DownloadButton(label="Download CSV", visible=False)
+
+    # This wrapper function handles the UI updates, which is the robust way to use Gradio.
+    def run_classification_and_update_ui(sem_image, cluster_thresh, m1_thresh, m2_thresh):
+        """
+        Calls the core logic and then returns updates for the Gradio UI components.
+        """
+        # Call the main processing function
+        classified_img, img_path, csv_path = damage_classification(sem_image, cluster_thresh, m1_thresh, m2_thresh)
+
+        # Return the results in the correct order to update the output components.
+        # Use gr.update to change properties of a component, like visibility.
+        return (
+            classified_img,
+            gr.update(value=img_path, visible=True),
+            gr.update(value=csv_path, visible=True)
+        )
+
+    # Connect the button's click event to the wrapper function
     classify_btn.click(
-        damage_classification,
-        inputs=[image_input, cluster_threshold_input, model1_threshold_input, model2_threshold_input],
-        outputs=[output_image, download_image_btn, download_csv_btn],
+        fn=run_classification_and_update_ui,
+        inputs=[
+            image_input,
+            cluster_threshold_input,
+            model1_threshold_input,
+            model2_threshold_input
+        ],
+        outputs=[
+            output_image,
+            download_image_btn,
+            download_csv_btn
+        ],
     )
 
 if __name__ == "__main__":
-    app.launch()
+    app.launch()