working example

.gitignore

@@ -188,4 +188,5 @@ myapp/debug_template.jpg
 
 debug_*.jpg
 *_clip_*.pth
-batch_run_*.txt
+batch_run_*.txt
+.claude/

classificator_training/helpers/args.py

@@ -56,7 +56,7 @@ def _override_args_from_model_name(args: Namespace, load_model_name: str, verbos
     )
 
     match = pattern.search(load_model_name)
-
+    print(match)
     if match:
         extracted_values = match.groupdict()
         

classificator_training/model/feature_extractor.py

@@ -162,4 +162,35 @@ class FeatureExtractor(nn.Module):
             extracted_feature_dict['vit'] = self.vit(self._ensure_batch_dim(vit_input))
             
         return extracted_feature_dict
+    
+    def extract_feature_from_model(self, model_name: str, input_tensor: torch.Tensor):
+        """
+        Extracts features from a specific backbone model.
+        """
+        if model_name not in self.use_models or not self.use_models[model_name]:
+            raise ValueError(f"Model '{model_name}' is not enabled in use_models.")
+        
+        input_tensor = self._ensure_batch_dim(input_tensor)
+        
+        if model_name == 'clip':
+            return self.clip_model(input_tensor).pooler_output
+        elif model_name == 'segformer':
+            output = self.segformer_model(input_tensor, output_hidden_states=True)
+            return output.last_hidden_state.mean(dim=[2, 3])
+        elif model_name == 'dpt':
+            output = self.dpt_model(input_tensor)
+            return output.pooler_output
+        elif model_name == 'midas':
+            output = self.midas_model(input_tensor, output_hidden_states=True)
+            return output.hidden_states[-1][:, 1:, :].mean(dim=1)
+        elif model_name == 'resnet':
+            return self.resnet(input_tensor)
+        elif model_name == 'mobilenet':
+            return self.mobilenet(input_tensor)
+        elif model_name == 'efficientnet':
+            return self.efficientnet(input_tensor)
+        elif model_name == 'vit':
+            return self.vit(input_tensor)
+        else:
+            raise ValueError(f"Unknown model name: {model_name}")
     

classificator_training/run_script.sh

@@ -185,8 +185,6 @@ RESULTS_FILE="batch_run_results2.txt"
 
 
 for SET in "${PARAMETER_SETS[@]}"; do
-    # Read 15 variables (plus P16/Alpha if needed, currently Alpha is treated as part of the set logic usually)
-    # Mapping based on new order:
     read P_CLIP P_SEG P_MIDAS P_DPT P_RES P_MOB P_EFF P_VIT P_GATE P_BATCH P_TYPE P_HEAD P_LR P_MARGIN P_ALPHA <<< "$SET"
 
     echo "--- Starting run with Models: Clip=$P_CLIP Seg=$P_SEG Midas=$P_MIDAS Dpt=$P_DPT Res=$P_RES Mob=$P_MOB Eff=$P_EFF Vit=$P_VIT ---"
@@ -210,7 +208,6 @@ for SET in "${PARAMETER_SETS[@]}"; do
 
     echo "$COMMAND_TO_RUN"
     
-    # Check if this is a "silent" run (P11 logic from your old script was unclear, assuming standard logging here)
     OUTPUT=$($COMMAND_TO_RUN 2>&1 | tee /dev/tty | tail -n 1)
     EXIT_CODE=${PIPESTATUS[0]}
 
@@ -222,7 +219,6 @@ for SET in "${PARAMETER_SETS[@]}"; do
         echo "Run finished **successfully**."
     else
         echo "Run **failed**. Stopping batch."
-        # break # Uncomment to stop on failure
     fi
     echo "---"
 done

myapp/.gitignore

@@ -171,4 +171,5 @@ cython_debug/
 .ruff_cache/
 
 # PyPI configuration file
-.pypirc
+.pypirc
+.claude/

myapp/AR/model_template.py

@@ -5,6 +5,6 @@ TRAINED_CLASSIFICATION_MODEL = {
     "feature_extractor": None,
     "id_to_tag": None,
     "tag_to_id": None,
-    "yolo_model": None,
+    "pose_model": None,
     "keypoints": None
 }

myapp/AR/pose_estimate_utils.py

@@ -1,8 +1,8 @@
-import json
 from typing import Dict, List, Tuple
-
 from transformers import CLIPConfig, CLIPProcessor, CLIPVisionModel
+from classificator_training.data.dataset import preprocess_single_image
 from classificator_training.model.feature_extractor import CLIP_MODEL_ID
+from classificator_training.utils import move_to_device
 from myapp.AR.model_template import TRAINED_CLASSIFICATION_MODEL
 import os
 import io
@@ -14,16 +14,9 @@ import numpy as np
 import torch
 import cv2
 import numpy as np
-import json
-import argparse
-from pathlib import Path
-from tqdm import tqdm
 import trimesh
 import pyrender
-from torchvision import transforms
-from datetime import datetime
 
-from myapp.AR.pose_network import BuildingPoseNet, quaternion_to_rotation_matrix
 
 def get_camera_matrix(img_w, img_h):
     """
@@ -41,183 +34,6 @@ def get_camera_matrix(img_w, img_h):
         [0, fy, cy],
         [0,  0,  1]
     ], dtype=np.float32)
-
-
-class BatchBuildingPoseEstimator:
-    """
-    Process folders of building images and generate pose predictions
-    Works with models trained on cached CLIP features
-    """
-    
-    def __init__(self, model_path, config_path, device='cuda'):
-        self.device = device
-        
-        # Load config
-        with open(config_path, 'r') as f:
-            self.config = json.load(f)
-        
-        print("Loading CLIP model for feature extraction...")
-        # Load CLIP for feature extraction
-        clip_config = CLIPConfig.from_pretrained(CLIP_MODEL_ID)
-        self.clip_model = CLIPVisionModel.from_pretrained(
-            CLIP_MODEL_ID,
-            config=clip_config.vision_config,
-        )
-        self.clip_model.to(device)
-        self.clip_model.eval()
-        
-        self.clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_ID)
-        
-        feature_dim = self.clip_model.config.hidden_size
-        print(f"✓ CLIP loaded (feature dim: {feature_dim})")
-        
-        # Load pose estimation model
-        print("Loading pose estimation model...")
-        self.model = BuildingPoseNet(
-            num_buildings=self.config['num_buildings'],
-            feature_dim=feature_dim
-        )
-        
-        checkpoint = torch.load(model_path, map_location=device, weights_only=False)
-        self.model.load_state_dict(checkpoint['model_state_dict'])
-        self.model.to(device)
-        self.model.eval()
-        
-        print(f"✓ Pose estimation model loaded successfully:")
-        print(f"  Model path: {model_path}")
-        print(f"  Device: {device}")
-        print(f"  Number of buildings: {self.config['num_buildings']}")
-    
-    def extract_clip_features(self, image):
-        """
-        Extract CLIP features from a single image
-        
-        Args:
-            image: numpy array (H, W, 3) in BGR or RGB format
-        
-        Returns:
-            features: torch.Tensor [1, feature_dim]
-        """
-        # Convert BGR to RGB if needed
-        if len(image.shape) == 2:
-            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-        elif image.shape[2] == 3:
-            # Assume BGR from OpenCV, convert to RGB
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        
-        # Process image with CLIP processor
-        inputs = self.clip_processor(images=image, return_tensors="pt")
-        pixel_values = inputs['pixel_values'].to(self.device)
-        
-        # Extract features
-        with torch.no_grad():
-            outputs = self.clip_model(pixel_values=pixel_values)
-            features = outputs.pooler_output  # [1, feature_dim]
-        
-        return features
-    
-    def predict_pose(self, image):
-        """
-        Predict pose from single image
-        
-        Args:
-            image: numpy array (H, W, 3) - raw image
-        
-        Returns:
-            dict with pose predictions
-        """
-        # Step 1: Extract CLIP features
-        features = self.extract_clip_features(image)
-        
-        # Step 2: Predict pose from features
-        with torch.no_grad():
-            output = self.model(features)
-        
-        # Parse outputs
-        building_id = torch.argmax(output['building_logits'], dim=1).item()
-        building_probs = torch.softmax(output['building_logits'], dim=1)[0].cpu().numpy()
-        rotation_quat = output['rotation'][0].cpu().numpy()
-        translation = output['translation'][0].cpu().numpy()
-        confidence = output['confidence'][0].item()
-        
-        # Convert quaternion to rotation matrix
-        rotation_matrix = quaternion_to_rotation_matrix(
-            torch.from_numpy(rotation_quat)
-        ).numpy()
-        
-        return {
-            'building_id': building_id,
-            'building_name': self.config['buildings'][building_id],
-            'building_probabilities': {
-                self.config['buildings'][i]: float(building_probs[i])
-                for i in range(len(building_probs))
-            },
-            'rotation_quaternion': rotation_quat.tolist(),
-            'rotation_matrix': rotation_matrix.tolist(),
-            'translation': translation.tolist(),
-            'confidence': float(confidence)
-        }
-    
-    def predict_pose_batch(self, images):
-        """
-        Predict poses for a batch of images (faster)
-        
-        Args:
-            images: list of numpy arrays
-        
-        Returns:
-            list of prediction dicts
-        """
-        # Convert all to RGB
-        rgb_images = []
-        for img in images:
-            if len(img.shape) == 2:
-                img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
-            elif img.shape[2] == 3:
-                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-            rgb_images.append(img)
-        
-        # Extract features for all images
-        inputs = self.clip_processor(images=rgb_images, return_tensors="pt")
-        pixel_values = inputs['pixel_values'].to(self.device)
-        
-        with torch.no_grad():
-            # Extract CLIP features
-            outputs = self.clip_model(pixel_values=pixel_values)
-            features = outputs.pooler_output  # [batch_size, feature_dim]
-            
-            # Predict poses
-            predictions = self.model(features)
-        
-        # Parse results
-        results = []
-        batch_size = len(images)
-        
-        for i in range(batch_size):
-            building_id = torch.argmax(predictions['building_logits'][i]).item()
-            building_probs = torch.softmax(predictions['building_logits'][i], dim=0).cpu().numpy()
-            rotation_quat = predictions['rotation'][i].cpu().numpy()
-            translation = predictions['translation'][i].cpu().numpy()
-            confidence = predictions['confidence'][i].item()
-            
-            rotation_matrix = quaternion_to_rotation_matrix(
-                torch.from_numpy(rotation_quat)
-            ).numpy()
-            
-            results.append({
-                'building_id': building_id,
-                'building_name': self.config['buildings'][building_id],
-                'building_probabilities': {
-                    self.config['buildings'][j]: float(building_probs[j])
-                    for j in range(len(building_probs))
-                },
-                'rotation_quaternion': rotation_quat.tolist(),
-                'rotation_matrix': rotation_matrix.tolist(),
-                'translation': translation.tolist(),
-                'confidence': float(confidence)
-            })
-        
-        return results
     
     
 def get_common_camera_intrinsics(img_width: int, img_height: int) -> List[Dict]:
@@ -439,13 +255,7 @@ def process_image_heuristic(img_bytes: bytes, model_bytes: bytes, model_name: st
         
     print("Running DL Pose Prediction...")
     
-    
-    pose_estimator = BatchBuildingPoseEstimator(
-        model_path='AR/best_model.pth',
-        config_path='AR/dataset.json',
-        device='cuda' if torch.cuda.is_available() else 'cpu',
-    )
-    
+    pose_estimator = TRAINED_CLASSIFICATION_MODEL['pose_model']
     pose_result = pose_estimator.predict_pose(img)
 
     R_blender = pose_result['rotation_matrix']
@@ -485,68 +295,6 @@ def process_image_heuristic(img_bytes: bytes, model_bytes: bytes, model_name: st
     rot_matrix[0:3, 0:3] = blender_to_opencv
     mesh_final.apply_transform(rot_matrix)
     
-    # camera_intrinsics = np.array([
-    #     [
-    #       796.4444444444445,
-    #       0,
-    #       512.0
-    #     ],
-    #     [
-    #       0,
-    #       796.4444444444445,
-    #       512.0
-    #     ],
-    #     [
-    #       0,
-    #       0,
-    #       1
-    #     ]
-    #   ])
-    
-    # fx = camera_intrinsics[0, 0]
-    # fy = camera_intrinsics[1, 1]
-    # cx = camera_intrinsics[0, 2]
-    # cy = camera_intrinsics[1, 2]
-    
-    # training_width = 1024
-    # training_height = 1024
-    # scale_x = w_img / training_width
-    # scale_y = h_img / training_height
-    
-    # fx_scaled = fx * scale_x
-    # fy_scaled = fy * scale_y
-    # cx_scaled = cx * scale_x
-    # cy_scaled = cy * scale_y
-
-    # view_matrix = np.eye(4)
-    # view_matrix[:3, :3] = R_opencv
-    # view_matrix[:3, 3] = t_opencv.squeeze()
-    
-    # cam_pose_cv = np.linalg.inv(view_matrix)
-    
-    # cv_to_gl = np.array([[1,0,0,0], [0,-1,0,0], [0,0,-1,0], [0,0,0,1]])
-    # final_pose = cam_pose_cv @ cv_to_gl
-
-    # scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0], ambient_light=[0.8, 0.8, 0.8])
-    # scene.add(pyrender.Mesh.from_trimesh(mesh_final))
-    
-    # cam = pyrender.IntrinsicsCamera(
-    #     fx=fx_scaled, 
-    #     fy=fy_scaled, 
-    #     cx=cx_scaled, 
-    #     cy=cy_scaled,
-    #     znear=0.05, 
-    #     zfar=1000.0
-    # )
-    
-    # scene.add(cam, pose=final_pose)
-    
-    # light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1000.0)
-    # scene.add(light, pose=final_pose)
-
-    # r = pyrender.OffscreenRenderer(w_img, h_img)
-    # color, depth = r.render(scene, flags=pyrender.RenderFlags.RGBA)
-    # r.delete()
     camera_configs = get_common_camera_intrinsics(w_img, h_img)[:3]
     
     best_score = -1

myapp/AR/pose_network.py

@@ -1,222 +1,12 @@
-# """
-# Pose Estimation Network for Building Alignment
-# Combines classification with pose regression for accurate 6DoF estimation
-# """
-
-# import torch
-# import torch.nn as nn
-# import torchvision.models as models
-# from torchvision.models import efficientnet_b3, EfficientNet_B3_Weights
-
-# class BuildingPoseNet(nn.Module):
-#     """
-#     Neural network for estimating 6DoF pose of buildings from images.
-    
-#     Architecture:
-#     - EfficientNet-B3 backbone (pre-trained on ImageNet)
-#     - Separate heads for rotation and translation
-#     - Rotation: quaternion representation (4 values)
-#     - Translation: 3D position (3 values)
-#     """
-    
-#     def __init__(self, num_buildings=10, pretrained=True):
-#         super(BuildingPoseNet, self).__init__()
-        
-#         if pretrained:
-#             weights = EfficientNet_B3_Weights.IMAGENET1K_V1
-#             self.backbone = efficientnet_b3(weights=weights)
-#         else:
-#             self.backbone = efficientnet_b3(weights=None)
-        
-#         feature_dim = self.backbone.classifier[1].in_features
-        
-#         self.backbone.classifier = nn.Identity()
-        
-#         self.building_classifier = nn.Sequential(
-#             nn.Linear(feature_dim, 512),
-#             nn.ReLU(),
-#             nn.Dropout(0.3),
-#             nn.Linear(512, num_buildings)
-#         )
-        
-#         self.rotation_head = nn.Sequential(
-#             nn.Linear(feature_dim, 512),
-#             nn.ReLU(),
-#             nn.Dropout(0.2),
-#             nn.Linear(512, 256),
-#             nn.ReLU(),
-#             nn.Linear(256, 4)
-#         )
-        
-#         self.translation_head = nn.Sequential(
-#             nn.Linear(feature_dim, 512),
-#             nn.ReLU(),
-#             nn.Dropout(0.2),
-#             nn.Linear(512, 256),
-#             nn.ReLU(),
-#             nn.Linear(256, 3)
-#         )
-        
-#         self.confidence_head = nn.Sequential(
-#             nn.Linear(feature_dim, 256),
-#             nn.ReLU(),
-#             nn.Linear(256, 1),
-#             nn.Sigmoid()
-#         )
-    
-#     def forward(self, x):
-#         features = self.backbone(x)
-        
-#         building_logits = self.building_classifier(features)
-#         rotation_quat = self.rotation_head(features)
-#         translation = self.translation_head(features)
-#         confidence = self.confidence_head(features)
-        
-#         rotation_quat = rotation_quat / (torch.norm(rotation_quat, dim=1, keepdim=True) + 1e-8)
-        
-#         return {
-#             'building_logits': building_logits,
-#             'rotation': rotation_quat,
-#             'translation': translation,
-#             'confidence': confidence
-#         }
-    
-
-# class PoseLoss(nn.Module):
-#     """
-#     Combined loss for pose estimation training
-#     """
-    
-#     def __init__(self, rotation_weight=1.0, translation_weight=1.0, 
-#                  classification_weight=0.5):
-#         super(PoseLoss, self).__init__()
-#         self.rotation_weight = rotation_weight
-#         self.translation_weight = translation_weight
-#         self.classification_weight = classification_weight
-#         self.ce_loss = nn.CrossEntropyLoss()
-        
-#     def quaternion_distance(self, q1, q2):
-#         """
-#         Compute geodesic distance between quaternions
-#         Returns angle in radians
-#         """
-#         q1 = q1 / (torch.norm(q1, dim=1, keepdim=True) + 1e-8)
-#         q2 = q2 / (torch.norm(q2, dim=1, keepdim=True) + 1e-8)
-        
-#         dot_product = torch.abs(torch.sum(q1 * q2, dim=1))
-#         dot_product = torch.clamp(dot_product, -1.0, 1.0)
-        
-#         return 2 * torch.acos(dot_product)
-    
-#     def forward(self, predictions, targets):
-#         """
-#         Args:
-#             predictions: dict with 'rotation', 'translation', 'building_logits'
-#             targets: dict with 'rotation', 'translation', 'building_id'
-#         """
-#         pred_rot = predictions['rotation']
-#         pred_rot = pred_rot / (torch.norm(pred_rot, dim=1, keepdim=True) + 1e-8)
-        
-#         target_rot = targets['rotation']
-#         target_rot = target_rot / (torch.norm(target_rot, dim=1, keepdim=True) + 1e-8)
-        
-#         rot_loss = self.quaternion_distance(pred_rot, target_rot).mean()
-        
-#         rot_loss = torch.clamp(rot_loss, 0, 10.0)
-        
-#         if torch.isnan(rot_loss):
-#             rot_loss = torch.tensor(1.0, device=rot_loss.device)
-        
-#         trans_loss = torch.nn.functional.mse_loss(
-#             predictions['translation'], 
-#             targets['translation']
-#         )
-        
-#         trans_loss = torch.clamp(trans_loss, 0, 100.0)
-
-#         if torch.isnan(trans_loss):
-#             trans_loss = torch.tensor(1.0, device=trans_loss.device)
-        
-#         cls_loss = 0
-#         if 'building_id' in targets and 'building_logits' in predictions:
-#             cls_loss = self.ce_loss(
-#                 predictions['building_logits'], 
-#                 targets['building_id']
-#             )
-#             if torch.isnan(cls_loss):
-#                 cls_loss = torch.tensor(0.1, device=predictions['building_logits'].device)
-        
-#         total_loss = (
-#             self.rotation_weight * rot_loss +
-#             self.translation_weight * trans_loss +
-#             self.classification_weight * cls_loss
-#         )
-        
-#         if torch.isnan(total_loss) or torch.isinf(total_loss):
-#             total_loss = torch.tensor(1.0, device=total_loss.device, requires_grad=True)
-        
-#         return {
-#             'total_loss': total_loss,
-#             'rotation_loss': rot_loss,
-#             'translation_loss': trans_loss,
-#             'classification_loss': cls_loss
-#         }
-
-# def quaternion_to_rotation_matrix(quaternion):
-#     """
-#     Convert quaternion to 3x3 rotation matrix
-    
-#     Args:
-#         quaternion: torch.Tensor of shape (4,) or (N, 4) [w, x, y, z]
-    
-#     Returns:
-#         Rotation matrix of shape (3, 3) or (N, 3, 3)
-#     """
-#     if quaternion.dim() == 1:
-#         quaternion = quaternion.unsqueeze(0)
-#         squeeze = True
-#     else:
-#         squeeze = False
-    
-#     quaternion = quaternion / torch.norm(quaternion, dim=1, keepdim=True)
-    
-#     w, x, y, z = quaternion[:, 0], quaternion[:, 1], quaternion[:, 2], quaternion[:, 3]
-    
-#     R = torch.zeros((quaternion.shape[0], 3, 3), device=quaternion.device)
-    
-#     R[:, 0, 0] = 1 - 2*y**2 - 2*z**2
-#     R[:, 0, 1] = 2*x*y - 2*w*z
-#     R[:, 0, 2] = 2*x*z + 2*w*y
-    
-#     R[:, 1, 0] = 2*x*y + 2*w*z
-#     R[:, 1, 1] = 1 - 2*x**2 - 2*z**2
-#     R[:, 1, 2] = 2*y*z - 2*w*x
-    
-#     R[:, 2, 0] = 2*x*z - 2*w*y
-#     R[:, 2, 1] = 2*y*z + 2*w*x
-#     R[:, 2, 2] = 1 - 2*x**2 - 2*y**2
-    
-#     if squeeze:
-#         R = R.squeeze(0)
-    
-#     return R
-
-# if __name__ == "__main__":
-#     model = BuildingPoseNet(num_buildings=5)
-    
-#     batch_size = 4
-#     dummy_input = torch.randn(batch_size, 3, 224, 224)
-    
-#     output = model(dummy_input)
-
-#     quat = output['rotation'][0]
-#     R = quaternion_to_rotation_matrix(quat)
-
-
-
+import json
+import cv2
 import torch
 import torch.nn as nn
 
+from classificator_training.data.dataset import preprocess_single_image
+from classificator_training.model.feature_extractor import FeatureExtractor
+from classificator_training.utils import move_to_device
+
 
 class BuildingPoseNet(nn.Module):
     """
@@ -443,32 +233,73 @@ def quaternion_to_rotation_matrix(quaternion):
     return R
 
 
-if __name__ == "__main__":
-    # Test the network
-    print("Testing BuildingPoseNetCached...")
-    
-    model = BuildingPoseNet(num_buildings=5, feature_dim=1280)
-    
-    # Create dummy features (simulating pre-extracted CLIP features)
-    batch_size = 4
-    dummy_features = torch.randn(batch_size, 1280)
-    
-    # Forward pass
-    output = model(dummy_features)
-    
-    print("\nNetwork Output Shapes:")
-    print(f"Building logits: {output['building_logits'].shape}")
-    print(f"Rotation (quaternion): {output['rotation'].shape}")
-    print(f"Translation: {output['translation'].shape}")
-    print(f"Confidence: {output['confidence'].shape}")
+
+class BuildingPoseEstimator:
+    """
+    Process folders of building images and generate pose predictions
+    Works with models trained on cached CLIP features
+    """
     
-    # Test quaternion to rotation matrix conversion
-    quat = output['rotation'][0]
-    R = quaternion_to_rotation_matrix(quat)
-    print(f"\nRotation matrix shape: {R.shape}")
-    print(f"Is orthogonal: {torch.allclose(R @ R.T, torch.eye(3), atol=1e-5)}")
+    def __init__(self, model_path, config_path, feature_extractor: FeatureExtractor, device='cuda'):
+        self.device = device
+        
+        with open(config_path, 'r') as f:
+            self.config = json.load(f)
+            
+        self.feature_extractor = feature_extractor
+        feature_dim = feature_extractor.feature_dims['clip']
+        
+        self.model = BuildingPoseNet(
+            num_buildings=self.config['num_buildings'],
+            feature_dim=feature_dim
+        )
+        
+        checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model.to(device)
+        self.model.eval()
     
-    # Count parameters
-    total_params = sum(p.numel() for p in model.parameters())
-    print(f"\nTotal parameters: {total_params:,}")
-    print("(Much smaller than full CLIP model!)")
+    def predict_pose(self, image):
+        """
+        Predict pose from single image
+        
+        Args:
+            image: numpy array (H, W, 3) - raw image
+        
+        Returns:
+            dict with pose predictions
+        """
+        # Step 1: Extract CLIP features
+        inputs = preprocess_single_image(image, self.feature_extractor.use_models)
+        inputs = move_to_device(inputs, "cuda" if torch.cuda.is_available() else "cpu")
+        features = self.feature_extractor.extract_feature_from_model("clip", inputs["clip_input"])
+        
+        # Step 2: Predict pose from features
+        with torch.no_grad():
+            output = self.model(features)
+        
+        # Parse outputs
+        building_id = torch.argmax(output['building_logits'], dim=1).item()
+        building_probs = torch.softmax(output['building_logits'], dim=1)[0].cpu().numpy()
+        rotation_quat = output['rotation'][0].cpu().numpy()
+        translation = output['translation'][0].cpu().numpy()
+        confidence = output['confidence'][0].item()
+        
+        # Convert quaternion to rotation matrix
+        rotation_matrix = quaternion_to_rotation_matrix(
+            torch.from_numpy(rotation_quat)
+        ).numpy()
+        
+        return {
+            'building_id': building_id,
+            'building_name': self.config['buildings'][building_id],
+            'building_probabilities': {
+                self.config['buildings'][i]: float(building_probs[i])
+                for i in range(len(building_probs))
+            },
+            'rotation_quaternion': rotation_quat.tolist(),
+            'rotation_matrix': rotation_matrix.tolist(),
+            'translation': translation.tolist(),
+            'confidence': float(confidence)
+        }
+        

myapp/main.py

@@ -3,20 +3,21 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from myapp.routers import buildings_router
 from myapp.routers import buildings_search_router
-from classificator_training.helpers.load_best_model import load_best_model
+from myapp.utils.load_best_model import load_best_model
 from myapp.AR.model_template import TRAINED_CLASSIFICATION_MODEL
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     print("Starting up: Loading Neural Network Model...")
     try:
-        model, prototype_tensor, class_ids, feature_extractor, id_to_tag, tag_to_id = load_best_model()
+        model, prototype_tensor, class_ids, feature_extractor, id_to_tag, tag_to_id, pose_model = load_best_model()
         TRAINED_CLASSIFICATION_MODEL["model"] = model
         TRAINED_CLASSIFICATION_MODEL["prototype_tensor"] = prototype_tensor
         TRAINED_CLASSIFICATION_MODEL["class_ids"] = class_ids
         TRAINED_CLASSIFICATION_MODEL["feature_extractor"] = feature_extractor
         TRAINED_CLASSIFICATION_MODEL["id_to_tag"] = id_to_tag
         TRAINED_CLASSIFICATION_MODEL['tag_to_id'] = tag_to_id
+        TRAINED_CLASSIFICATION_MODEL["pose_model"] = pose_model
         TRAINED_CLASSIFICATION_MODEL["model"].eval()
 
         print("Model loaded.")

{classificator_training/helpers → myapp/utils}/load_best_model.py

@@ -1,11 +1,11 @@
 from argparse import Namespace
-import json
 import os
+from pathlib import Path
 import torch
 from classificator_training.model.model import FusedFeatureModel
 from classificator_training.model.feature_extractor import FeatureExtractor
 from classificator_training.helpers.args import _override_args_from_model_name
-from ultralytics import YOLO
+from myapp.AR.pose_network import BuildingPoseEstimator
 
 ENVIRONEMENT = os.getenv("ENVIRONMENT", "production")
 MODEL_PATH = "1_fused_feature_model.pth_full_clip1_segformer0_midas0_dpt0_gate0_batch64_traintypehardmining_bigfusionhead2_lr2e-07_margin1.2_alpha64.0_datasetsize114272_rendersBlenderRenders7_testdatatest_data3.model"
@@ -18,9 +18,9 @@ FULL_YOLO_MODEL_PATH = FULL_MODEL_PATH + YOLO_MODEL_PATH
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def load_best_model():
-    args = Namespace(clip=1, segformer=0, midas=0, dpt=0, gate=0, big_fusion_head=2, train_type='hardmining', lr=2e-07, margin=0.8, alpha=64.0)
+    args = Namespace(clip=1, segformer=0, midas=0, dpt=0, gate=0, big_fusion_head=2, train_type='hardmining', lr=2e-07, margin=1.2, alpha=64.0)
     args = _override_args_from_model_name(args, MODEL_PATH)
-
+    print("Loading model with args:", args)
     MODELS_USED = {
         'clip': args.clip,
         'segformer': args.segformer,
@@ -41,7 +41,17 @@ def load_best_model():
     class_ids = checkpoint.get('class_ids', None)
     id_to_tag = checkpoint.get('id_to_tag', None)
     tag_to_id = {tag: id for id, tag in id_to_tag.items()}
-
-    return model, prototype_tensor, class_ids, feature_extractor, id_to_tag, tag_to_id
+    
+    current_dir = Path(__file__).parent 
+    file_path = current_dir.parent / 'AR'
+    print("Loaded")
+    pose_model = BuildingPoseEstimator(
+        model_path=file_path / 'best_model.pth',
+        config_path=file_path / 'dataset.json',
+        feature_extractor=feature_extractor,
+        device='cuda' if torch.cuda.is_available() else 'cpu',
+    )
+
+    return model, prototype_tensor, class_ids, feature_extractor, id_to_tag, tag_to_id, pose_model