Update app.py

app.py

@@ -1,35 +1,153 @@
 import gradio as gr
 import cv2
 import numpy as np
-from stitching import Stitcher
+from typing import Union, List
+from pathlib import Path
 from PIL import Image
+import torch
+from ultralytics.utils.plotting import save_one_box
+from ultralytics.engine.results import Results
+from ultralytics import YOLOWorld
 
-# Function to stitch uploaded images
-def stitch_images(images):
-    stitcher = Stitcher()
-    panorama = stitcher.stitch(images)  # Stitch the images
-    
-    # Convert the result to a PIL image
-    pil_image = Image.fromarray(cv2.cvtColor(panorama, cv2.COLOR_BGR2RGB))
-    
-    # Save the image in JPG format
-    pil_image.save("stitched_image.jpg", "JPEG")
-    
-    # Return the stitched image
-    return pil_image, "stitched_image.jpg"  # Return the image and the file path
-
-# Create Gradio interface with button
+# Function to resize images
+
+def resize_images(images, scale_percent=50):
+    resized_images = []
+    for img in images:
+        width = int(img.shape[1] * scale_percent / 100)
+        height = int(img.shape[0] * scale_percent / 100)
+        dim = (width, height)
+        resized = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
+        resized_images.append(resized)
+    return resized_images
+
+# Function to stitch images
+
+def stitch_images(image_paths, scale_percent=50):
+    images = [cv2.imread(path) for path in image_paths]
+    resized_images = resize_images(images, scale_percent)
+    stitcher = cv2.Stitcher_create()
+    status, stitched_image = stitcher.stitch(resized_images)
+
+    if status == cv2.Stitcher_OK:
+        print("Stitching successful!")
+        return stitched_image
+    else:
+        print(f"Stitching failed with status code: {status}")
+        return None
+
+# YOLO detection function
+
+def yolo_detect(
+    image: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor],
+    classes: List[str],
+) -> np.ndarray:
+    detector = YOLOWorld("yolov8x-worldv2.pt")
+    detector.set_classes(classes)
+    results: List[Results] = detector.predict(image)
+
+    boxes = []
+    for result in results:
+        boxes.append(
+            save_one_box(result.cpu().boxes.xyxy, im=result.orig_img, save=False)
+        )
+
+    del detector
+    return boxes[0]
+
+# Function to predict and process image
+
+def predict(image, offset_inches):
+    try:
+        drawer_img = yolo_detect(image, ["box"])
+        shrunked_img = make_square(shrink_bbox(drawer_img, 0.8))
+        return shrunked_img
+    except:
+        raise Exception("Unable to DETECT DRAWER, please take another picture with different magnification level!")
+
+# Function to shrink bounding box
+
+def shrink_bbox(image: np.ndarray, shrink_factor: float):
+    height, width = image.shape[:2]
+    center_x, center_y = width // 2, height // 2
+
+    new_width = int(width * shrink_factor)
+    new_height = int(height * shrink_factor)
+
+    x1 = max(center_x - new_width // 2, 0)
+    y1 = max(center_y - new_height // 2, 0)
+    x2 = min(center_x + new_width // 2, width)
+    y2 = min(center_y + new_height // 2, height)
+
+    cropped_image = image[y1:y2, x1:x2]
+    return cropped_image
+
+# Function to make image square
+
+def make_square(img: np.ndarray):
+    height, width = img.shape[:2]
+    max_dim = max(height, width)
+    pad_height = (max_dim - height) // 2
+    pad_width = (max_dim - width) // 2
+    pad_height_extra = max_dim - height - 2 * pad_height
+    pad_width_extra = max_dim - width - 2 * pad_width
+
+    if len(img.shape) == 3:
+        padded = np.pad(
+            img,
+            ((pad_height, pad_height + pad_height_extra),
+             (pad_width, pad_width + pad_width_extra),
+             (0, 0)),
+            mode="edge"
+        )
+    else:
+        padded = np.pad(
+            img,
+            ((pad_height, pad_height + pad_height_extra),
+             (pad_width, pad_width + pad_width_extra)),
+            mode="edge"
+        )
+
+    return padded
+
+# Main image processing function
+
+def process_image(image_paths, scale_percent=50, offset_inches=1):
+    stitched_image = stitch_images(image_paths, scale_percent)
+
+    if stitched_image is not None:
+        try:
+            stitched_image_rgb = cv2.cvtColor(stitched_image, cv2.COLOR_BGR2RGB)
+            final_image = predict(stitched_image_rgb, offset_inches)
+            return final_image
+        except Exception as e:
+            print(str(e))
+            return stitched_image
+
+# Gradio interface function
+
+def gradio_stitch_and_detect(image_files):
+    image_paths = [file.name for file in image_files]
+    result_image = process_image(image_paths, scale_percent=50)
+
+    if result_image is not None:
+        result_image_rgb = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(result_image_rgb)
+        pil_image.save("stitched_image.jpg", "JPEG")
+        return pil_image, "stitched_image.jpg"
+
+    return None, None
+
+# Gradio interface
 with gr.Blocks() as interface:
-    gr.Markdown("<h1 style='color: #2196F3; text-align: center;'>Image Stitcher 🧵</h1>")  # Heading with a new color
-    gr.Markdown("<h3 style='color: #2196F3; text-align: center;'>=== Upload the images you want to stitch ===</h3>")  # Subheading
-    
-    image_upload = gr.Files(type="filepath", label="Upload Images")  # File input for images
-    stitch_button = gr.Button("Stitch", variant="primary")  # Button with color
-    stitched_image = gr.Image(type="pil", label="Stitched Image")  # Display stitched image
-    download_button = gr.File(label="Download Stitched Image")  # Button to download the image
-
-    # Define button interaction
-    stitch_button.click(stitch_images, inputs=image_upload, outputs=[stitched_image, download_button])
-
-# Launch the interface
-interface.launch()
+    gr.Markdown("<h1 style='color: #2196F3; text-align: center;'>Image Stitcher 🧵</h1>")
+    gr.Markdown("<h3 style='color: #2196F3; text-align: center;'>=== Upload the images you want to stitch ===</h3>")
+
+    image_upload = gr.Files(type="filepath", label="Upload Images")
+    stitch_button = gr.Button("Stitch", variant="primary")
+    stitched_image = gr.Image(type="pil", label="Stitched Image")
+    download_button = gr.File(label="Download Stitched Image")
+
+    stitch_button.click(gradio_stitch_and_detect, inputs=image_upload, outputs=[stitched_image, download_button])
+
+interface.launch()