use size before resize to normalize Euc distance in WHD

Former-commit-id: c6b28714559909ef9879efe7f713bc106af64acf
javiribera · Mar 29, 2018 · 2ed82ff · 2ed82ff
1 parent c23d015
commit 2ed82ff
Show file tree

Hide file tree

Showing 5 changed files with 441 additions and 20 deletions.
diff --git a/object-locator/data.py b/object-locator/data.py
@@ -11,6 +11,7 @@
 from torchvision import transforms
 import xmltodict
 from parse import parse
+from . import get_image_size
 
 
 class CSVDataset(data.Dataset):
@@ -221,6 +222,10 @@ def __call__(self, img, dictionary):
                                              xs.view(-1, 1)),
                                             1)
 
+        # Indicate new size in dictionary
+        dictionary['resized_height'] = self.size[0]
+        dictionary['resized_width'] = self.size[1]
+
         return img, dictionary
 
 
@@ -346,8 +351,13 @@ def __init__(self,
                         locations = []
                         for plant in plot['plants']['plant']:
                             locations.append(eval(plant['location_wrt_plot']))
+                        img_abspath = os.path.join(self.root_dir, filename)
+                        orig_width, orig_height = \
+                            get_image_size.get_image_size(img_abspath)
                         self.dict[filename] = {'count': count,
-                                               'locations': locations}
+                                               'locations': locations,
+                                               'orig_width': orig_width,
+                                               'orig_height': orig_height}
 
             # Use an Ordered Dictionary to allow random access
             self.dict = OrderedDict(self.dict.items())
@@ -380,7 +390,11 @@ def __getitem__(self, idx):
             filename, dictionary = self.dict_list[idx]
         else:
             filename = self.listfiles[idx]
-            dictionary = {'filename': self.listfiles[idx]}
+            orig_width, orig_height = \
+                get_image_size.get_image_size(img_abspath)
+            dictionary = {'filename': self.listfiles[idx],
+                          'orig_width': orig_width,
+                          'orig_height': orig_height}
         img_abspath = os.path.join(self.root_dir, filename)
 
         img = Image.open(img_abspath)