Update app.py

app.py

@@ -112,9 +112,10 @@ class RAGSystem:
             return ""
 
 class ImageAnalyzer:
-    def __init__(self):
-        self.device = "cpu"  # Force CPU usage for better compatibility
+    def __init__(self, model_name="microsoft/swin-base-patch4-window7-224-in22k"):
+        self.device = "cpu"
         self.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
+        self.model_name = model_name
         self._model = None
         self._feature_extractor = None
 
@@ -127,22 +128,60 @@ class ImageAnalyzer:
     @property
     def feature_extractor(self):
         if self._feature_extractor is None:
-            self._feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+            self._feature_extractor = self._load_feature_extractor()
         return self._feature_extractor
 
+    def _load_feature_extractor(self):
+        """Load the appropriate feature extractor based on model type"""
+        try:
+            if "swin" in self.model_name:
+                from transformers import AutoFeatureExtractor
+                return AutoFeatureExtractor.from_pretrained(self.model_name)
+            elif "convnext" in self.model_name:
+                from transformers import ConvNextFeatureExtractor
+                return ConvNextFeatureExtractor.from_pretrained(self.model_name)
+            else:
+                from transformers import ViTFeatureExtractor
+                return ViTFeatureExtractor.from_pretrained(self.model_name)
+        except Exception as e:
+            logger.error(f"Feature extractor initialization error: {e}")
+            return None
+
     def _load_model(self):
         try:
-            model = ViTForImageClassification.from_pretrained(
-                "google/vit-base-patch16-224",
-                num_labels=len(self.defect_classes),
-                ignore_mismatched_sizes=True
-            ).to(self.device)
+            if "swin" in self.model_name:
+                from transformers import SwinForImageClassification
+                model = SwinForImageClassification.from_pretrained(
+                    self.model_name,
+                    num_labels=len(self.defect_classes),
+                    ignore_mismatched_sizes=True
+                )
+            elif "convnext" in self.model_name:
+                from transformers import ConvNextForImageClassification
+                model = ConvNextForImageClassification.from_pretrained(
+                    self.model_name,
+                    num_labels=len(self.defect_classes),
+                    ignore_mismatched_sizes=True
+                )
+            else:
+                from transformers import ViTForImageClassification
+                model = ViTForImageClassification.from_pretrained(
+                    self.model_name,
+                    num_labels=len(self.defect_classes),
+                    ignore_mismatched_sizes=True
+                )
+
+            model = model.to(self.device)
             
+            # Reinitialize the classifier layer
             with torch.no_grad():
-                model.classifier = torch.nn.Linear(
-                    in_features=model.classifier.in_features,
-                    out_features=len(self.defect_classes)
-                )
+                if hasattr(model, 'classifier'):
+                    in_features = model.classifier.in_features
+                    model.classifier = torch.nn.Linear(in_features, len(self.defect_classes))
+                elif hasattr(model, 'head'):
+                    in_features = model.head.in_features
+                    model.head = torch.nn.Linear(in_features, len(self.defect_classes))
+                
             return model
         except Exception as e:
             logger.error(f"Model initialization error: {e}")
@@ -150,7 +189,7 @@ class ImageAnalyzer:
 
     def preprocess_image(self, image_bytes):
         """Preprocess image for model input"""
-        return _cached_preprocess_image(image_bytes)
+        return _cached_preprocess_image(image_bytes, self.model_name)
 
     def analyze_image(self, image):
         """Analyze image for defects"""
@@ -187,20 +226,26 @@ class ImageAnalyzer:
             return None
 
 @st.cache_data
-def _cached_preprocess_image(image_bytes):
+def _cached_preprocess_image(image_bytes, model_name):
     """Cached version of image preprocessing"""
     try:
         image = Image.open(image_bytes)
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
-        width, height = 224, 224
+        # Adjust size based on model requirements
+        if "convnext" in model_name:
+            width, height = 384, 384
+        else:
+            width, height = 224, 224
+            
         image = image.resize((width, height), Image.Resampling.LANCZOS)
         return image
     except Exception as e:
         logger.error(f"Image preprocessing error: {e}")
         return None
-        
+
+@st.cache_data      
 def get_groq_response(query: str, context: str) -> str:
     """Get response from Groq LLM with caching"""
     try: