Fixes to get inference and training working with reorganized files

common/detector.py

@@ -38,7 +38,6 @@
         # transforms.Scale(IMAGE_SIZE),
         # transforms.CenterCrop((imagesize,imagesize)),
         transforms.ToTensor(),
-        # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
         transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
     ]
 )

common/utils.py

@@ -1,37 +1,25 @@
 """
 NVIDIA from [email protected]
 """
+import albumentations as A
+import colorsys
+import glob
+import io
+import json
+import math
+from math import acos
+from math import sqrt
+from math import pi
 import numpy as np
-import torch
-
 import os
-
+from os.path import exists, basename, join
+from PIL import Image, ImageDraw, ImageEnhance, ImageFont
 import torch
 import torch.nn as nn
 import torch.nn.parallel
-
-import torch.utils.data
-
-import torchvision.transforms as transforms
-
 import torch.utils.data as data
-import glob
-import os
-import io
-
-from PIL import Image
-from PIL import ImageDraw
-from PIL import ImageEnhance
-
-from math import acos
-from math import sqrt
-from math import pi
-
-from os.path import exists, basename
-import json
-from os.path import join
+import torchvision.transforms as transforms
 
-import albumentations as A
 
 
 def default_loader(path):
@@ -74,8 +62,6 @@ def py_ang(A, B=(1, 0)):
         return 360 - inner
 
 
-import colorsys, math
-
 
 def append_dot(extensions):
     res = []
@@ -858,7 +844,6 @@ def save_image(tensor, filename, nrow=4, padding=2, mean=None, std=None, save=Tr
     Saves a given Tensor into an image file.
     If given a mini-batch tensor, will save the tensor as a grid of images.
     """
-    from PIL import Image
 
     tensor = tensor.cpu()
     grid = make_grid(tensor, nrow=nrow, padding=10, pad_value=1)
@@ -889,10 +874,6 @@ def save_image(tensor, filename, nrow=4, padding=2, mean=None, std=None, save=Tr
     return im, grid
 
 
-from PIL import ImageDraw, Image, ImageFont
-import json
-
-
 class Draw(object):
     """Drawing helper class to visualize the neural network output"""
 
@@ -943,3 +924,18 @@ def draw_cube(self, points, color=(0, 255, 0)):
             self.draw_text(points[i], str(i), (255, 0, 0))
 
 
+def get_image_grid(tensor, nrow=3, padding=2, mean=None, std=None):
+    """
+    Saves a given Tensor into an image file.
+    If given a mini-batch tensor, will save the tensor as a grid of images.
+    """
+
+    # tensor = tensor.cpu()
+    grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=1)
+    if not mean is None:
+        # ndarr = grid.mul(std).add(mean).mul(255).byte().transpose(0,2).transpose(0,1).numpy()
+        ndarr = grid.mul(std).add(mean).mul(255).byte().transpose(0, 2).transpose(0, 1).numpy()
+    else:
+        ndarr = grid.mul(0.5).add(0.5).mul(255).byte().transpose(0, 2).transpose(0, 1).numpy()
+    im = Image.fromarray(ndarr)
+    return im

inference/README.md

@@ -22,7 +22,7 @@ The `inference.py` script will take a trained model to run inference. In order t
 Below is an example of running inference:
 
 ```
-python inference.py --weights ../output/weights --data ../sample_data --object 003_cracker_box
+python inference.py --weights ../output/weights --data ../sample_data --object cracker
 ```
 
 ### Configuration Files

inference/inference.py

@@ -1,19 +1,18 @@
 #!/usr/bin/env python
 
+import argparse
 import cv2
 import numpy as np
 from PIL import Image
-
+import os
+import simplejson as json
 import sys
+import yaml
 
 sys.path.append("../common/")
-
 from cuboid import Cuboid3d
 from cuboid_pnp_solver import CuboidPNPSolver
 from detector import ModelData, ObjectDetector
-
-import simplejson as json
-
 from utils import loadimages_inference, loadweights, Draw
 
 
@@ -156,11 +155,6 @@ def image_callback(
 
 
 if __name__ == "__main__":
-
-    import argparse
-    import yaml
-    import os
-
     parser = argparse.ArgumentParser()
 
     parser.add_argument(
@@ -175,12 +169,12 @@ def image_callback(
     )
     parser.add_argument(
         "--config",
-        default="config/config_pose.yaml",
+        default="../config/config_pose.yaml",
         help="Path to inference config file",
     )
     parser.add_argument(
         "--camera",
-        default="config/camera_info.yaml",
+        default="../config/camera_info.yaml",
         help="Path to camera info file",
     )
 
@@ -226,7 +220,7 @@ def image_callback(
         exit()
     else:
         print(f"Found {len(weights)} weights. ")
-    # Load inference images
+        # Load inference images
     imgs, imgsname = loadimages_inference(opt.data, extensions=opt.exts)
 
     if len(imgs) == 0 or len(imgsname) == 0:
@@ -235,7 +229,7 @@ def image_callback(
         )
         exit()
 
-    for w_i, weight in enumerate(weights):
+    for w_i, weight in enumerate(weights): 
         dope_node = DopeNode(config, weight, opt.object)
 
         for i in range(len(imgs)):

requirements.txt

@@ -1,4 +1,4 @@
-torch==2.1.2
+torch>=2.1.0
 torchvision
 tensorboardX
 boto3

train/sample_data/000000.json → sample_data/000000.json

@@ -2,7 +2,7 @@
   "camera_data": {},
   "objects": [
     {
-      "class": "003_cracker_box",
+      "class": "cracker",
       "visibility": 1.0,
       "location": [
         -5.40710973739624,
@@ -55,4 +55,4 @@
       ]
     }
   ]
-}
+}

train/sample_data/000001.json → sample_data/000001.json

@@ -2,7 +2,7 @@
   "camera_data": {},
   "objects": [
     {
-      "class": "003_cracker_box",
+      "class": "cracker",
       "visibility": 1.0,
       "location": [
         -5.40710973739624,
@@ -55,4 +55,4 @@
       ]
     }
   ]
-}
+}

train/sample_data/000030.json → sample_data/000030.json

@@ -2,7 +2,7 @@
   "camera_data": {},
   "objects": [
     {
-      "class": "003_cracker_box",
+      "class": "cracker",
       "visibility": 1.0,
       "location": [
         6.001536846160889,
@@ -55,4 +55,4 @@
       ]
     }
   ]
-}
+}

train/sample_data/000031.json → sample_data/000031.json

@@ -2,7 +2,7 @@
   "camera_data": {},
   "objects": [
     {
-      "class": "003_cracker_box",
+      "class": "cracker",
       "visibility": 1.0,
       "location": [
         6.001536846160889,
@@ -55,4 +55,4 @@
       ]
     }
   ]
-}
+}