fix: Add error raise for not found model weight in inference. ruff style update

Author: yakhyo

config.py

@@ -1,15 +1,12 @@
 data_config = {
-    "gaze360":
-        {
-            "bins": 90,
-            "binwidth": 4,
-            "angle": 180  # angle range
-        },
-    "mpiigaze":
-        {
-            "bins": 28,
-            "binwidth": 3,
-            "angle": 42  # angle range
-        }
-
-}
+    "gaze360": {
+        "bins": 90,
+        "binwidth": 4,
+        "angle": 180,  # angle range
+    },
+    "mpiigaze": {
+        "bins": 28,
+        "binwidth": 3,
+        "angle": 42,  # angle range
+    },
+}

evaluate.py

@@ -13,25 +13,40 @@
 from utils.helpers import angular_error, gaze_to_3d, get_dataloader, get_model
 
 import warnings
+
 warnings.filterwarnings("ignore")
 # Setup logging
-logging.basicConfig(level=logging.INFO, format='%(message)s')
+logging.basicConfig(level=logging.INFO, format="%(message)s")
 
 
 def parse_args():
     """Parse input arguments."""
     parser = argparse.ArgumentParser(description="Gaze estimation evaluation")
-    parser.add_argument("--data", type=str, default="data/Gaze360", help="Directory path for gaze images.")
-    parser.add_argument("--dataset", type=str, default="gaze360", help="Dataset name, available `gaze360`, `mpiigaze`")
-    parser.add_argument("--weight", type=str, default="", help="Path to model weight for evaluation.")
+    parser.add_argument(
+        "--data",
+        type=str,
+        default="data/Gaze360",
+        help="Directory path for gaze images.",
+    )
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="gaze360",
+        help="Dataset name, available `gaze360`, `mpiigaze`",
+    )
+    parser.add_argument(
+        "--weight", type=str, default="", help="Path to model weight for evaluation."
+    )
     parser.add_argument("--batch-size", type=int, default=64, help="Batch size.")
     parser.add_argument(
         "--arch",
         type=str,
         default="resnet18",
-        help="Network architecture, currently available: resnet18/34/50, mobilenetv2, mobileone_s0-s4."
+        help="Network architecture, currently available: resnet18/34/50, mobilenetv2, mobileone_s0-s4.",
+    )
+    parser.add_argument(
+        "--num-workers", type=int, default=8, help="Number of workers for data loading."
     )
-    parser.add_argument("--num-workers", type=int, default=8, help="Number of workers for data loading.")
 
     args = parser.parse_args()
 
@@ -42,7 +57,9 @@ def parse_args():
         args.binwidth = dataset_config["binwidth"]
         args.angle = dataset_config["angle"]
     else:
-        raise ValueError(f"Unknown dataset: {args.dataset}. Available options: {list(data_config.keys())}")
+        raise ValueError(
+            f"Unknown dataset: {args.dataset}. Available options: {list(data_config.keys())}"
+        )
 
     return args
 
@@ -63,7 +80,9 @@ def evaluate(params, model, data_loader, idx_tensor, device):
     average_error = 0
     total_samples = 0
 
-    for images, labels_gaze, regression_labels_gaze, _ in tqdm(data_loader, total=len(data_loader)):
+    for images, labels_gaze, regression_labels_gaze, _ in tqdm(
+        data_loader, total=len(data_loader)
+    ):
         total_samples += regression_labels_gaze.size(0)
         images = images.to(device)
 
@@ -79,8 +98,12 @@ def evaluate(params, model, data_loader, idx_tensor, device):
         yaw_predicted = F.softmax(yaw, dim=1)
 
         # Mapping from binned (0 to 90) to angles (-180 to 180) or (0 to 28) to angles (-42, 42)
-        pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1) * params.binwidth - params.angle
-        yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1) * params.binwidth - params.angle
+        pitch_predicted = (
+            torch.sum(pitch_predicted * idx_tensor, 1) * params.binwidth - params.angle
+        )
+        yaw_predicted = (
+            torch.sum(yaw_predicted * idx_tensor, 1) * params.binwidth - params.angle
+        )
 
         pitch_predicted = np.radians(pitch_predicted.cpu())
         yaw_predicted = np.radians(yaw_predicted.cpu())
@@ -91,7 +114,7 @@ def evaluate(params, model, data_loader, idx_tensor, device):
     logging.info(
         f"Dataset: {params.dataset} | "
         f"Total Number of Samples: {total_samples} | "
-        f"Mean Angular Error: {average_error/total_samples}"
+        f"Mean Angular Error: {average_error / total_samples}"
     )
 
 
@@ -104,7 +127,9 @@ def main():
     model = get_model(params.arch, params.bins, inference_mode=True)
 
     if os.path.exists(params.weight):
-        model.load_state_dict(torch.load(params.weight, map_location=device, weights_only=True))
+        model.load_state_dict(
+            torch.load(params.weight, map_location=device, weights_only=True)
+        )
     else:
         raise ValueError(f"Model weight not found at {params.weight}")
 
@@ -117,5 +142,5 @@ def main():
     evaluate(params, model, test_loader, idx_tensor, device)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

inference.py

@@ -14,23 +14,41 @@
 from uniface import RetinaFace
 
 warnings.filterwarnings("ignore")
-logging.basicConfig(level=logging.INFO, format='%(message)s')
+logging.basicConfig(level=logging.INFO, format="%(message)s")
 
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Gaze estimation inference")
-    parser.add_argument("--model", type=str, default="resnet34", help="Model name, default `resnet18`")
+    parser.add_argument(
+        "--model", type=str, default="resnet34", help="Model name, default `resnet18`"
+    )
     parser.add_argument(
         "--weight",
         type=str,
         default="resnet34.pt",
-        help="Path to gaze esimation model weights"
+        help="Path to gaze esimation model weights",
+    )
+    parser.add_argument(
+        "--view",
+        action="store_true",
+        default=True,
+        help="Display the inference results",
+    )
+    parser.add_argument(
+        "--source",
+        type=str,
+        default="assets/in_video.mp4",
+        help="Path to source video file or camera index",
+    )
+    parser.add_argument(
+        "--output", type=str, default="output.mp4", help="Path to save output file"
+    )
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="gaze360",
+        help="Dataset name to get dataset related configs",
     )
-    parser.add_argument("--view", action="store_true", default=True, help="Display the inference results")
-    parser.add_argument("--source", type=str, default="assets/in_video.mp4",
-                        help="Path to source video file or camera index")
-    parser.add_argument("--output", type=str, default="output.mp4", help="Path to save output file")
-    parser.add_argument("--dataset", type=str, default="gaze360", help="Dataset name to get dataset related configs")
     args = parser.parse_args()
 
     # Override default values based on selected dataset
@@ -40,19 +58,23 @@ def parse_args():
         args.binwidth = dataset_config["binwidth"]
         args.angle = dataset_config["angle"]
     else:
-        raise ValueError(f"Unknown dataset: {args.dataset}. Available options: {list(data_config.keys())}")
+        raise ValueError(
+            f"Unknown dataset: {args.dataset}. Available options: {list(data_config.keys())}"
+        )
 
     return args
 
 
 def pre_process(image):
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    transform = transforms.Compose([
-        transforms.ToPILImage(),
-        transforms.Resize(448),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
+    transform = transforms.Compose(
+        [
+            transforms.ToPILImage(),
+            transforms.Resize(448),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
 
     image = transform(image)
     image_batch = image.unsqueeze(0)
@@ -72,13 +94,16 @@ def main(params):
         gaze_detector.load_state_dict(state_dict)
         logging.info("Gaze Estimation model weights loaded.")
     except Exception as e:
-        logging.info(f"Exception occured while loading pre-trained weights of gaze estimation model. Exception: {e}")
+        logging.info(
+            f"Exception occured while loading pre-trained weights of gaze estimation model. Exception: {e}"
+        )
+        raise FileNotFoundError(f"Model weights not found at {params.weight}") from e
 
     gaze_detector.to(device)
     gaze_detector.eval()
 
     video_source = params.source
-    if video_source.isdigit() or video_source == '0':
+    if video_source.isdigit() or video_source == "0":
         cap = cv2.VideoCapture(int(video_source))
     else:
         cap = cv2.VideoCapture(video_source)
@@ -87,7 +112,9 @@ def main(params):
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        out = cv2.VideoWriter(params.output, fourcc, cap.get(cv2.CAP_PROP_FPS), (width, height))
+        out = cv2.VideoWriter(
+            params.output, fourcc, cap.get(cv2.CAP_PROP_FPS), (width, height)
+        )
 
     if not cap.isOpened():
         raise IOError("Cannot open webcam")
@@ -102,7 +129,7 @@ def main(params):
 
             faces = face_detector.detect(frame)
             for face in faces:
-                bbox = face['bbox']
+                bbox = face["bbox"]
                 x_min, y_min, x_max, y_max = map(int, bbox[:4])
 
                 image = frame[y_min:y_max, x_min:x_max]
@@ -111,11 +138,20 @@ def main(params):
 
                 pitch, yaw = gaze_detector(image)
 
-                pitch_predicted, yaw_predicted = F.softmax(pitch, dim=1), F.softmax(yaw, dim=1)
+                pitch_predicted, yaw_predicted = (
+                    F.softmax(pitch, dim=1),
+                    F.softmax(yaw, dim=1),
+                )
 
                 # Mapping from binned (0 to 90) to angles (-180 to 180) or (0 to 28) to angles (-42, 42)
-                pitch_predicted = torch.sum(pitch_predicted * idx_tensor, dim=1) * params.binwidth - params.angle
-                yaw_predicted = torch.sum(yaw_predicted * idx_tensor, dim=1) * params.binwidth - params.angle
+                pitch_predicted = (
+                    torch.sum(pitch_predicted * idx_tensor, dim=1) * params.binwidth
+                    - params.angle
+                )
+                yaw_predicted = (
+                    torch.sum(yaw_predicted * idx_tensor, dim=1) * params.binwidth
+                    - params.angle
+                )
 
                 # Degrees to Radians
                 pitch_predicted = np.radians(pitch_predicted.cpu())
@@ -128,8 +164,8 @@ def main(params):
                 out.write(frame)
 
             if params.view:
-                cv2.imshow('Demo', frame)
-                if cv2.waitKey(1) & 0xFF == ord('q'):
+                cv2.imshow("Demo", frame)
+                if cv2.waitKey(1) & 0xFF == ord("q"):
                     break
 
     cap.release()