new

app.py

@@ -8,6 +8,14 @@ import plotly.graph_objects as go
 import torch
 import spaces
 
+import os
+import shutil
+import subprocess
+
+import os
+import shutil
+import subprocess
+
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 
@@ -15,7 +23,7 @@ class Detection:
 
     def __init__(self):
         # Set the model path and confidence threshold
-        yolov8_model_path = "./model/best.pt"  # Update to your model path
+        yolov8_model_path = "./model/best_100epochs_latest.pt"  # Update to your model path
 
         # Initialize the AutoDetectionModel
         self.model = AutoDetectionModel.from_pretrained(
@@ -97,7 +105,7 @@ class Detection:
         # Populate the category_areas dictionary
         for annotation in annotations:
             category_name = annotation['category_name']
-            area = annotation['bbox'][2] * annotation['bbox'][3]  # Width * Height
+            area = annotation['area']
             if category_name in category_areas:
                 category_areas[category_name].append(area)
 
@@ -299,6 +307,54 @@ def generate_confidence_bar_chart(annotations):
 
 
 
+# Directory to save images
+img_dir = "./stitching/img_dir/"
+output_dir = "./"
+os.makedirs(img_dir, exist_ok=True)
+os.makedirs(output_dir, exist_ok=True)
+
+# Function to handle the stitching process
+def save_and_stitch(first_image, second_image, third_image, fourth_image):
+    # Save images to `img_dir`
+    images = [first_image, second_image, third_image, fourth_image]
+    for idx, img in enumerate(images):
+        if img is not None:
+            file_path = os.path.join(img_dir, f"Image_{idx + 1}.jpg")
+            img.save(file_path, format="JPEG")
+
+    # Execute the stitching command for all image files in the folder
+    command = f"stitch {img_dir}/Image_*.jpg"
+    try:
+        subprocess.run(command, shell=True, check=True)
+
+        # Load the result image from ./stitching/result.jpg
+        result_image_path = os.path.join(output_dir, "result.jpg")
+        if os.path.exists(result_image_path):
+            print("found")
+            return Image.open(result_image_path)
+        else:
+            print("not found")
+            return None  # If result image doesn't exist, return None
+
+    except subprocess.CalledProcessError as e:
+        print(f"Error executing command: {str(e)}")
+        return None
+
+# Function to clear the img_dir
+def clear_img_dir():
+    for file_name in os.listdir(img_dir):
+        file_path = os.path.join(img_dir, file_name)
+        try:
+            if os.path.isfile(file_path):
+                os.remove(file_path)
+            elif os.path.isdir(file_path):
+                os.rmdir(file_path)  # For directories, remove them
+        except Exception as e:
+            print(f"Error deleting file {file_name}: {str(e)}")
+    return "Images cleared from img_dir!"
+
+
+
 
 
 
@@ -335,11 +391,11 @@ with gr.Blocks() as demo:
             }
                 
 
-            .gradio-container-5-9-0 .prose * {
+            .gradio-container-5-4-0 .prose * {
                 color: #083484;
             }
 
-            .gradio-container-5-9-0 .prose :first-child {
+            .gradio-container-5-4-0 .prose :first-child {
                 margin-top: 85px
             }
 
@@ -421,7 +477,7 @@ with gr.Blocks() as demo:
                 font-family: 'Montserrat',sans-serif;
             }   
                 
-            .gradio-container-5-9-0 .prose h1, .gradio-container-5-9-0 .prose h2, .gradio-container-5-9-0 .prose h3, .gradio-container-5-9-0 .prose h4, .gradio-container-5-9-0 .prose h5 {
+            .gradio-container-5-4-0 .prose h1, .gradio-container-5-4-0 .prose h2, .gradio-container-5-4-0 .prose h3, .gradio-container-5-4-0 .prose h4, .gradio-container-5-4-0 .prose h5 {
             margin: var(--spacing-xxl) 0 var(--spacing-lg);
             font-weight: var(--prose-header-text-weight);
             line-height: 1.3;
@@ -517,7 +573,7 @@ with gr.Blocks() as demo:
         pittings_confidence_bar_chart = gr.Plot(label="Pittings Confidence Score Distribution")
 
 
-    
+    #Gradio row for confidence bar chart
     with gr.Row(visible=False) as confidence_btn_row:
         confidence_chart_btn = gr.Button("Generate Confidence Chart", variant="primary")
         confidence_chart_btn.click(
@@ -525,6 +581,54 @@ with gr.Blocks() as demo:
             inputs=output_annotations,  # Pass the annotations
             outputs=[nicks_confidence_bar_chart,dents_confidence_bar_chart,scratches_confidence_bar_chart,pittings_confidence_bar_chart]
         )
+
+
+
+
+
+    with gr.Row(visible=False) as upload_image_stitching:
+        first_image_stitching = gr.Image(type="pil", label="Select Image 1")
+        second_image_stitching = gr.Image(type="pil", label="Select Image 2")
+        third_image_stitching = gr.Image(type="pil", label="Select Image 3")
+        fourth_image_stitching = gr.Image(type="pil", label="Select Image 4")
+
+    # Row for result output
+    with gr.Row(visible=False) as result_output_block:
+        result_output = gr.Image(type="pil",label="Stitched Output Image")
+    
+    # Row for buttons
+    with gr.Row(visible=False) as stitching_btn:
+        apply_stitching_btn = gr.Button("Apply Stitching",variant="primary")
+        # Button click actions
+        apply_stitching_btn.click(
+            save_and_stitch,
+            inputs=[
+                first_image_stitching,
+                second_image_stitching,
+                third_image_stitching,
+                fourth_image_stitching,
+            ],
+            outputs=result_output,
+        )
+
+
+    # Row for displaying status
+    with gr.Row(visible=False) as display_img_dir_status:
+        status_text = gr.Textbox(label="Status")
+        
+        
+    # Row for clearing images from img_dir
+    with gr.Row(visible=False) as clear_img_btn:
+        clear_img_dir_btn = gr.Button("Clear Images from img_dir",variant="primary")
+        clear_img_dir_btn.click(
+        clear_img_dir,
+        inputs=[],
+        outputs=status_text
+    )
+
+
+    
+
     
      # Login row, initially visible
     with gr.Row(visible=True) as login_row:
@@ -553,6 +657,11 @@ with gr.Blocks() as demo:
                             gr.update(visible=login_state),  # Show frequency_btn_row
                             gr.update(visible=login_state),  #Show Confidence chart
                             gr.update(visible=login_state),  #Show Confidence btn
+                            gr.update(visible=login_state),  #Show upload image stitching
+                            gr.update(visible=login_state),  #Show stitched result output
+                            gr.update(visible=login_state),  #Show stitching btn
+                            gr.update(visible=login_state),  #Show display image dir status
+                            gr.update(visible=login_state),  #Show clear image btn
                             gr.update(visible=not login_state) # for login
                         ),
                         inputs=login_successful,
@@ -564,6 +673,11 @@ with gr.Blocks() as demo:
                                  frequency_btn_row,
                                  confidence_bar_chart_row, 
                                  confidence_btn_row, 
+                                 upload_image_stitching,
+                                 result_output_block,
+                                 stitching_btn,
+                                 display_img_dir_status,
+                                 clear_img_btn,
                                  login_row]
                     )
 

requirements.txt

@@ -1,3 +1,4 @@
+git+https://github.com/OpenStitching/stitching.git@main#egg=stitching
 --extra-index-url https://download.pytorch.org/whl/cu124
 torch
 torchvision