Merge pull request #667 from hesamossanloo/main

Update data type for bounding box coordinates to float32
weecology · May 11, 2024 · dfd9f31 · dfd9f31
2 parents a03a481 + 0f12089
commit dfd9f31
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/deepforest/dataset.py b/deepforest/dataset.py
@@ -100,7 +100,7 @@ def __getitem__(self, idx):
                                                  self.image_names[idx]]
             targets = {}
             targets["boxes"] = image_annotations[["xmin", "ymin", "xmax",
-                                                  "ymax"]].values.astype(float)
+                                                  "ymax"]].values.astype("float32")
 
             # Labels need to be encoded
             targets["labels"] = image_annotations.label.apply(
@@ -112,7 +112,7 @@ def __getitem__(self, idx):
                 labels = torch.from_numpy(targets["labels"])
                 # channels last
                 image = np.rollaxis(image, 2, 0)
-                image = torch.from_numpy(image)
+                image = torch.from_numpy(image).float()
                 targets = {"boxes": boxes, "labels": labels}
                 return self.image_names[idx], image, targets
 
@@ -122,7 +122,7 @@ def __getitem__(self, idx):
             image = augmented["image"]
 
             boxes = np.array(augmented["bboxes"])
-            boxes = torch.from_numpy(boxes)
+            boxes = torch.from_numpy(boxes).float()
             labels = np.array(augmented["category_ids"])
             labels = torch.from_numpy(labels)
             targets = {"boxes": boxes, "labels": labels}