first commit

LINEA

@@ -0,0 +1 @@
+Subproject commit 12ac0a326ddb7ec9809bf7080a19c8509dcffe45

app.py

@@ -0,0 +1,225 @@
+import gradio as gr
+import torch
+from torch import nn
+import torchvision.transforms as T
+
+from LINEA.models import build_linea
+from LINEA.util.slconfig import DictAction, SLConfig
+
+from PIL import Image, ImageDraw
+
+LINEA_MODELS = {
+    "LINEA-N": './LINEA/configs/linea/linea_hgnetv2_n.py',
+    "LINEA-S": './LINEA/configs/linea/linea_hgnetv2_s.py',
+    "LINEA-M": './LINEA/configs/linea/linea_hgnetv2_m.py',
+    "LINEA-L": './LINEA/configs/linea/linea_hgnetv2_l.py'
+}
+
+transforms = T.Compose(
+        [
+            T.Resize((640, 640)),
+            T.ToTensor(),
+            T.Normalize(mean=[0.538, 0.494, 0.453], std=[0.257, 0.263, 0.273]),
+        ]
+    )
+
+example_images = [
+    ["assets/example1.jpg"],
+    ["assets/example2.jpg"],
+    ["assets/example3.jpg"],
+    ["assets/example4.jpg"],
+]
+
+description = """
+<h1 align="center">
+  <ins>LINEA</ins>
+  <br>
+  Fast and accurate line detection using scalable transformers
+</h1>
+
+<h2 align="center">
+<a href="https://www.linkedin.com/in/sebastianjr/">Sebastian Janampa</a> 
+and 
+<a href="https://www.linkedin.com/in/marios-pattichis-207b0119/">Marios Pattichis</a>
+</h2>
+
+<h2 align="center">  
+    <a href="https://github.com/SebastianJanampa/LINEA.git">GitHub</a> |
+    <a href="https://colab.research.google.com/github/SebastianJanampa/LINEA/blob/master/LINEA_tutorial.ipynb">Colab</a>
+</h2>
+
+
+## Getting Started
+
+LINEA is a family of transformers models that detectes the line segments on an image. 
+Its key component is its new attention mechanism called **line attention**.
+
+To get started, upload an image or select one of the examples below. 
+You can choose between different model size, change the confidence threshold and visualize the results.
+"""
+
+def create_model(model_size):
+    cfg = SLConfig.fromfile(LINEA_MODELS[model_size])
+    cfg.pretrained = False
+
+    model, postprocessor = build_linea(cfg)
+
+    letter = model_size[-1].lower()
+    url = f"https://github.com/SebastianJanampa/storage/releases/download/LINEA/linea_hgnetv2_{letter}.pth"
+    state_dict = torch.hub.load_state_dict_from_url(
+            url, map_location="cpu", file_name=f"linea_hgnetv2_{letter}.pth"
+        )
+
+    model.load_state_dict(state_dict['model'], strict=True)
+
+    class Model(nn.Module):
+      def __init__(self):
+        super().__init__()
+        self.model = model.deploy()
+        self.postprocessor = postprocessor.deploy()
+
+      def forward(self, images, orig_target_sizes):
+        outputs = self.model(images)
+        outputs = self.postprocessor(outputs, orig_target_sizes)
+        return outputs
+
+    model = Model()
+    model.eval()
+
+    return model
+
+def draw(images, lines, scores, thrh):
+    for i, im in enumerate(images):
+        draw = ImageDraw.Draw(im)
+
+        scr = scores[i]
+        line = lines[i][scr > thrh]
+        scrs = scr[scr > thrh]
+
+        for j, l in enumerate(line):
+            draw.line(list(l), fill="red", width=5)
+            draw.text(
+                (l[0], l[1]),
+                text=f"{round(scrs[j].item(), 2)}",
+                fill="blue",
+            )
+
+    return images
+
+def filter(lines, scores, threshold):
+    filtered_lines, filter_scores = [], []
+    for line, scr in zip(lines, scores):
+        idx = scr > threshold
+        filtered_lines.append(line[idx])
+        filter_scores.append(scr[idx])
+    return filtered_lines, filter_scores
+
+def format_output(lines, scores):
+    n = len(lines[0])
+
+    txt = f"{n} lines were detected\n"
+    txt += "Detected lines:\n"
+    for line, scr in zip(lines[0], scores[0]):
+        txt += f"\tx1: {line[0].item():.2f}"
+        txt += f"\ty1: {line[0].item():.2f}"
+        txt += f"\tx2: {line[0].item():.2f}"
+        txt += f"\ty2: {line[0].item():.2f}"
+        txt += f"\tscore: {scr.item():.2f}\n"
+    return txt
+
+def process_results(
+    image_path,
+    model_size, 
+    threshold
+    ):
+    """ Process the image an returns the detected lines """
+    if image_path is None:
+        raise gr.Error("Please upload an image first.")
+
+    model = create_model(model_size)
+
+    im_pil = Image.open(image_path).convert("RGB")
+    w, h = im_pil.size
+    orig_size = torch.tensor([[w, h]])
+
+    im_data = transforms(im_pil).unsqueeze(0)
+
+    output = model(im_data, orig_size)
+    lines, scores = output
+
+    result_images = draw([im_pil], lines, scores, thrh=threshold)
+    filtered_lines, filtered_scores = filter(lines, scores, threshold)
+
+    return format_output(filtered_lines, filtered_scores), result_images[0], (lines, scores)
+
+def update_threshold(
+    image_path, 
+    raw_results,
+    threshold
+    ):
+    lines, scores = raw_results
+    im_pil = Image.open(image_path).convert("RGB")
+
+    result_images = draw([im_pil], lines, scores, thrh=threshold)
+    filtered_lines, filtered_scores = filter(lines, scores, threshold)
+    return format_output(filtered_lines, filtered_scores), result_images[0]
+
+def update_model(
+    image_path,
+    model_size, 
+    threshold
+    ):
+    create_model(model_size)
+
+    if image_path is None:
+        raise gr.Error("Please upload an image first.")
+        return None, None, None
+
+    return process_results(image_path, model_size, threshold)
+
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown(description)
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("""## Input Image""")
+            image_path = gr.Image(label="Upload image", type="filepath")
+            model_size = gr.Dropdown(
+                choices=list(LINEA_MODELS.keys()), label="Choose a LINEA model.", value="LINEA-M"
+            )
+            threshold = gr.Slider(
+                label="Confidence Threshold",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.05,
+                interactive=True,
+                value=0.30,
+            )
+
+            submit_btn = gr.Button("Detect Lines")
+            gr.Examples(examples=example_images, inputs=[image_path, model_size])
+
+        with gr.Column():
+            gr.Markdown("""## Results""")
+            image_output = gr.Image(label="Detected Lines")
+
+            text_output = gr.Textbox(label="Predicted lines", type="text", lines=5)
+
+    # Define the action when the button is clicked
+    raw_results = gr.State()
+
+    plot_inputs = [
+        raw_results,
+        threshold
+    ]
+
+    submit_btn.click(
+        fn=process_results,
+        inputs=[image_path, model_size] + plot_inputs[1:],
+        outputs=[text_output, image_output, raw_results],
+    )
+
+    # Define the action when the plot checkboxes are clicked
+    threshold.change(fn=update_threshold, inputs=[image_path] + plot_inputs, outputs=[text_output, image_output])
+demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+torch>=2.0.1
+torchvision>=0.15.2
+transformers
+yapf
+addict
+scipy