create app.py

app.py

@@ -0,0 +1,390 @@
+"""
+Gradio web application for detecting and measuring objects in images.
+
+Key features:
+- Image scaling tool to set measurement reference
+- Object detection using YOLOv8 model for scallop/spat detection
+- Interactive annotation of detected objects
+- Size measurements in mm based on scale reference
+- Statistics and histogram visualization of object sizes
+- Export results to CSV
+
+## TODO:
+- [ ] Load annotations from T-Rex
+
+"""
+
+# %% #|> Imports											 |
+from pathlib import Path
+import cv2
+import gradio as gr
+from gradio_image_annotation import image_annotator
+import numpy as np
+import pandas as pd
+import supervision as sv
+
+import plotly.express as px
+
+from spatstatapp.inference import inference_large
+from spatstatapp.plotting import coco_to_detections
+from spatstatapp.tile_training_data import load_bboxes
+
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw
+import cv2
+
+# %%	load image and detections										 |
+model_path = Path("models/best.pt")
+
+data_dir=Path("img")
+data_dir.exists()
+
+train_images = list(data_dir.glob('shells.png'))
+
+default_images = {Path(img).stem: str(img) for img in train_images}
+
+class PointSelector:
+    def __init__(self, image=None):
+        self.points = []
+        self.og_img = image
+        self.image_path = image
+        self.line_len_px = None
+        self.line_len_mm = None
+
+    def reset(self):
+        self.points = []
+        return self.og_img, "Points cleared"
+
+    def reset_og_img(self, image):
+        # self.og_img = image.copy()
+        # raise Exception(image)
+        self.og_img = None
+        self.points = []
+        self.image_path = image
+
+    def add_point(self, image, evt: gr.SelectData):
+        img_draw = cv2.imread(image)#[:,:,::-1]
+        if (len(self.points) == 0):# & (self.og_img is None):
+            self.image_path = image
+        #     self.og_img = cv2.imread(image)
+        #     img_draw = self.og_img.copy()
+
+        if len(self.points) >= 2:
+            self.points = []
+            img_draw = cv2.imread(self.image_path)
+            # img_draw = self.og_img.copy()
+
+        self.points.append((evt.index[0], evt.index[1]))
+
+        # Draw on image
+        # img_draw = image.copy()
+        if len(self.points) > 0:
+            for pt in self.points:
+                cv2.circle(img_draw, (int(pt[0]), int(pt[1])), 5, (255,0,0), -1)
+
+        if len(self.points) == 2:
+            cv2.line(img_draw,
+                    (int(self.points[0][0]), int(self.points[0][1])),
+                    (int(self.points[1][0]), int(self.points[1][1])),
+                    (0,255,0), 3)
+
+            # Calculate distance
+            dist = np.sqrt((self.points[1][0] - self.points[0][0])**2 +
+                         (self.points[1][1] - self.points[0][1])**2)
+            msg = f"Distance: {dist:.1f} pixels"
+            self.line_len_px = dist
+        else:
+            msg = f"Click point {len(self.points)+1}"
+
+        return img_draw[:,:,::-1], msg
+
+    def set_line_length(self, line_len_mm, button):
+        self.line_len_mm = line_len_mm
+        return self.check_scale_set(button)
+
+    def check_scale_set(self, button):
+        if (self.line_len_mm is not None) & (self.line_len_px is not None):
+        # if True:
+            return gr.update(visible=True)
+        else:
+            return gr.update(visible=False)
+
+    def save_scaled_boxes(self, annotator):
+        try:
+            json_data = annotator["boxes"]
+            if len(json_data)==0:
+                return None
+            else:
+                df = pd.DataFrame(json_data).drop(columns=["color"], errors='ignore')
+                df["xrange"] = ((df["xmax"] - df["xmin"])*(self.line_len_mm/self.line_len_px)).round(2)
+                df["yrange"] = ((df["ymax"] - df["ymin"])*(self.line_len_mm/self.line_len_px)).round(2)
+                df["mean_daimeter_mm"] = ((df["yrange"]+df["xrange"])/2).round(2)
+
+                return df
+        except Exception as e:
+            return None
+
+
+
+
+def detections_to_json(detections:sv.Detections, image:np.ndarray):
+    """Add predictions to canvas"""
+    boxes = []
+
+    for xyxy, mask, confidence, class_id, tracker_id, data in detections:
+        xmin, ymin, xmax, ymax = xyxy
+        obj = {
+            "xmin": float(xmin),
+            "ymin": float(ymin),
+            "xmax": float(xmax),
+            "ymax": float(ymax),
+            "label": "",# data["class_name"],
+            "color": (255, 0, 0)
+        }
+        boxes.append(obj)
+
+    annotation = {
+        "image": image,
+        "boxes": boxes
+            }
+
+    return annotation
+
+
+
+
+def create_histogram(df):
+    # print(type(df))
+    # print(len(df))
+    print()
+    if df is None or len(df) == 0 or df.iloc[0,0]=="":
+        return None
+    fig = px.histogram(df, x="mean_daimeter_mm",
+                      title="Distribution of Shell Sizes",
+                      labels={"mean_daimeter_mm": "Mean Diameter (mm)"},
+                      nbins=30)
+    return fig
+
+
+# def get_boxes_table(annotator):
+#     json_data = annotator["boxes"]
+#     if len(json_data)==0:
+#         return pd.DataFrame()
+#     else:
+#         df = pd.DataFrame(json_data).drop(columns=["color"], errors='ignore')
+#         return df
+from ultralytics.utils.ops import xywhn2xyxy
+
+
+
+def find_boxes_json(image_path):
+    # print(annotator)
+    img = cv2.imread(image_path)
+    detections = inference_large(img, model_path, sam_path=None, edge_pct=0.01, conf_threshold=0.4, overlap_px=100, tile_px=400)
+    annotations = detections_to_json(detections, image_path)
+    annotations["image"] = image_path
+    # annotator.update(annotations)
+    annotator = image_annotator(
+            annotations,
+            boxes_alpha=0.02,
+            handle_size=4,
+            show_label=False,
+        )
+    return annotator, annotations["boxes"]
+
+def load_coco_boxes(image_path, coco_file, class_labels="scallop"):
+    image = cv2.imread(image_path)[:,:,::-1]
+    detections = coco_to_detections(coco_file, image)
+    annotations = detections_to_json(detections, image)
+    annotations["image"] = image_path
+    # annotator.update(annotations)
+    annotator = image_annotator(
+            annotations,
+            boxes_alpha=0.02,
+            handle_size=4,
+            show_label=False,
+        )
+    return annotator, annotations["boxes"]
+
+
+selector = PointSelector()
+
+# %%	Tab 1										 |
+with gr.Blocks() as demo:
+    with gr.Tabs() as tabs:
+        with gr.TabItem("Scale ", id=0):
+            with gr.Row():
+                with gr.Column(scale=10):
+                    image_input = gr.Image(label="Click two points to measure distance", type="filepath", interactive=True)
+
+                    default_image = gr.Dropdown(
+                        choices=["None"] + list(default_images.keys()),
+                        label="Use default image?",
+                        # value=list(default_images.keys())[0] if default_images else None
+                    )
+                with gr.Column(scale=1, min_width=200):
+                    filename = gr.Textbox(label="Filename")
+                    output_text = gr.Textbox(label="Status", value="Click two points to measure distance")
+                    line_length_mm = gr.Number(label="line length in mm")
+                    target_select = gr.Radio(label ="select target:", visible=True, choices=["scallop", "spat"])
+                    button_find = gr.Button("find bounding boxes", visible=False)
+                    load_annot_btn = gr.Button("Load Existing boxes (Optional)",  visible=False)
+                    load_annot = gr.File(label="Load Existing boxes (Optional)",file_types=[".txt"], file_count="single", visible=False, height=500)
+                    # test_text = gr.Textbox(label="test")
+
+
+
+# %%	T1: event handlers										 |
+                # image_input.upload()
+            default_image.change(
+                lambda x: default_images[x] if x in default_images.keys() else None,
+                inputs=[default_image],
+                outputs=[image_input]
+            )
+
+            image_input.upload(
+                selector.reset_og_img,
+                inputs=[image_input],
+            )
+            image_input.upload(
+                lambda x: Path(x).name,
+                inputs=[image_input],
+                outputs=[filename]
+            )
+
+            # Event handlers
+            image_input.select(
+                selector.add_point,
+                inputs=[image_input],
+                outputs=[image_input, output_text]
+            )
+
+            line_length_mm.change(
+                selector.set_line_length,
+                inputs=[line_length_mm, button_find],
+                outputs=[button_find]
+            )
+
+            line_length_mm.change(
+                selector.check_scale_set,
+                inputs = load_annot_btn,
+                outputs=load_annot_btn,
+            )
+
+            load_annot_btn.click(
+                lambda: gr.update(visible=True),
+                outputs=[load_annot]
+            )
+
+            # load_annot.upload(
+            #     load_bboxes,
+            #     inputs=[load_annot],
+            #     outputs=[test_text]
+            # )
+
+# %%	Tab2										 |
+        with gr.TabItem("Object annotation", id=1, visible=True):
+            annotator = image_annotator(
+                boxes_alpha=0.02,
+                handle_size=4,
+                show_label=False,
+                label_list=["scallop", "spat"],
+                label_colors=[(255, 0, 0), (255, 200, 0)]
+            )
+
+            # button_get = gr.Button("Get bounding boxes")
+            download_file = gr.File(
+                label="Download CSV",
+                visible=True,
+                # interactive=True
+            )
+            with gr.Row():
+                with gr.Column(scale=1):
+                    obj_count = gr.Textbox(label="Object count")
+                    # button_save = gr.Button("save bounding boxes")
+                with gr.Column(scale=1):
+                    obj_size = gr.Textbox(value = "Has the scale size been set?" ,label="Mean size")
+
+            histogram = gr.Plot()
+            table = gr.DataFrame(
+                max_height=500,
+            )
+            # table = gr.Textbox(label="Status", value=1)
+
+            json_data = gr.JSON(value={}, visible=False)
+
+
+# %%	T2: event handlers										 |
+            json_boxes = button_find.click(
+                fn=find_boxes_json,
+                inputs=[image_input],
+                outputs=[annotator, json_data]
+            )
+
+            button_find.click(
+                fn=lambda: gr.Tabs(selected=1),
+                outputs=tabs
+            )
+
+            json_boxes = load_annot.upload(
+                fn=load_coco_boxes,
+                inputs=[image_input, load_annot],
+                outputs=[annotator, json_data]
+            )
+
+            # button_find.click(
+            #     fn=change_tab,
+            #     # inputs=[annotator, image_input],
+            #     outputs=tabs
+            # )
+
+            # annotator.change(
+            # json_boxes = button_get.click(
+            json_boxes = annotator.change(
+                fn=selector.save_scaled_boxes,
+                inputs= [annotator],
+                outputs= table
+                )
+
+            table.change(
+                fn=create_histogram,
+                inputs=[table],
+                outputs=[histogram]
+            )
+
+            def df_mean_count(df):
+                try:
+                    mean = df["mean_daimeter_mm"].mean().round(2)
+                    count = len(df)
+                    return mean, count
+                except Exception as e:
+                    return "Has the scale size been set?", None
+
+            table.change(
+                fn=df_mean_count,
+                inputs=[table],
+                outputs=[obj_size, obj_count]
+            )
+
+            def save_and_download_table(df, img_name):
+                try:
+                    # Create temporary file with .csv extension
+                    # with NamedTemporaryFile(delete=False, suffix='.csv') as tmp_file:
+                    #     csv_path = tmp_file.name
+                    csv_path = Path(img_name).stem +"_boxes.csv"
+                    df.to_csv(csv_path, index=False)
+                    return csv_path
+                except Exception as e:
+                    return None
+
+            table.change(
+                    fn=save_and_download_table,
+                    inputs=[table, filename],
+                    outputs=[download_file]
+                )
+
+
+if __name__ == "__main__":
+    demo.launch()
+