Update floorplan.py

HanHan55 · Sep 28, 2021 · 8a6fee8 · 8a6fee8
1 parent def9083
commit 8a6fee8
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/Network/model/floorplan.py b/Network/model/floorplan.py
@@ -155,7 +155,13 @@ def get_boxes(self,relative=True,tensor=True):
             X,Y = np.linspace(0,1,w),np.linspace(0,1,h)
             boxes = boxes-np.array([y0,x0,y0,x0])
 
-        norm = lambda box:np.array([X[box[1]],Y[box[0]],X[box[3]-1],Y[box[2]-1]])
+        norm = lambda box:np.array([
+            X[max(box[1],0)],
+            Y[max(box[0],0)],
+            X[min(box[3]-1,w-1)],
+            Y[min(box[2]-1,h-1)]
+        ])
+
         boxes = np.apply_along_axis(norm,1,boxes)
         boxes[:,2:]-=boxes[:,:2]
         boxes[:,:2]+=boxes[:,2:]/2
@@ -347,4 +353,4 @@ def vis_fp(fp,size=(256,256)):
     return img
 
 if __name__ == "__main__":
-    pass
+    pass