fix: resolve IndexError in line fitting visualization

app.py

@@ -1,13 +1,10 @@
-# Example showing how to fit a 2d line with kornia / pytorch
 import matplotlib.pyplot as plt
 import torch
 import matplotlib
 matplotlib.use('Agg')
-import matplotlib.pyplot as plt
 import gradio as gr
 from kornia.geometry.line import ParametrizedLine, fit_line
 
-
 def inference(point1, point2, point3, point4):
     std = 1.2  # standard deviation for the points
     num_points = 50  # total number of points
@@ -16,7 +13,7 @@ def inference(point1, point2, point3, point4):
     p0 = torch.tensor([point1, point2], dtype=torch.float32)
     p1 = torch.tensor([point3, point4], dtype=torch.float32)
     l1 = ParametrizedLine.through(p0, p1)
-    
+
     # sample some points and weights
     pts, w = [], []
     for t in torch.linspace(-10, 10, num_points):
@@ -25,37 +22,41 @@ def inference(point1, point2, point3, point4):
         p2 += p2_noise
         pts.append(p2)
         w.append(1 - p2_noise.mean())
+    
     pts = torch.stack(pts)
     w = torch.stack(w)
+
     if len(pts.shape) == 2:
-        pts = pts[None]
+        pts = pts.unsqueeze(0)
     if len(w.shape) == 1:
-        w = w[None]
-    
+        w = w.unsqueeze(0)
+
     l2 = fit_line(pts, w)
-    
-    # project some points along the estimated line
-    p3 = l2.point_at(-10)
-    p4 = l2.point_at(10)
 
-    X = torch.stack((p3, p4)).detach().numpy()
-    X_pts = pts.detach().numpy()
+    # project some points along the estimated line
+    p3 = l2.point_at(torch.tensor(-10.0))
+    p4 = l2.point_at(torch.tensor(10.0))
+    X = torch.stack((p3, p4)).squeeze().detach().numpy()
+    X_pts = pts.squeeze().detach().numpy()
 
-    fig = plt.figure()
-    plt.plot(X_pts[:, 0], X_pts[:, 1], 'ro')
-    plt.plot(X[:, 0], X[:, 1])
+    fig, ax = plt.subplots()
+    ax.plot(X_pts[:, 0], X_pts[:, 1], 'ro')
+    ax.plot(X[:, 0], X[:, 1])
+    ax.set_xlim(X_pts[:, 0].min() - 1, X_pts[:, 0].max() + 1)
+    ax.set_ylim(X_pts[:, 1].min() - 1, X_pts[:, 1].max() + 1)
     return fig
 
 inputs = [
-    gr.inputs.Slider(0.0, 10.0, default=0.0, label="Point 1"),
-    gr.inputs.Slider(0.0, 10.0, default=0.0, label="Point 2"),
-    gr.inputs.Slider(0.0, 10.0, default=0.0, label="Point 3"),
-    gr.inputs.Slider(0.0, 10.0, default=0.0, label="Point 4"),
+    gr.Slider(0.0, 10.0, value=0.0, label="Point 1 X"),
+    gr.Slider(0.0, 10.0, value=0.0, label="Point 1 Y"),
+    gr.Slider(0.0, 10.0, value=10.0, label="Point 2 X"),
+    gr.Slider(0.0, 10.0, value=10.0, label="Point 2 Y"),
 ]
+
 outputs = gr.Plot()
 
 examples = [
-    [0.0, 0.0, 1.0, 1.0],
+    [0.0, 0.0, 10.0, 10.0],
 ]
 
 title = 'Line Fitting'
@@ -69,6 +70,6 @@ demo = gr.Interface(
     theme='huggingface',
     live=True,
     examples=examples,
-
 )
-demo.launch()
+
+demo.launch()