Update app.py

app.py

@@ -17,22 +17,49 @@ import os
 import json
 from pathlib import Path
 
-# Import the model architecture
-import sys
-import os
+# Define the model architecture directly (standalone)
+import torch
+import torch.nn as nn
+from torchvision import models
+
+class MultiOutputModel(nn.Module):
+    """Multi-output model for artifact classification"""
+
+    def __init__(self, num_object_classes, num_material_classes, hidden_size=512):
+        super(MultiOutputModel, self).__init__()
+
+        # Use a pre-trained ResNet as backbone
+        self.backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
+        # Remove the final classification layer
+        self.backbone = nn.Sequential(*list(self.backbone.children())[:-1])
+
+        # Freeze early layers for transfer learning
+        for param in list(self.backbone.parameters())[:-2]:
+            param.requires_grad = False
+
+        # Classification heads for each attribute
+        self.object_classifier = nn.Linear(2048, num_object_classes)
+        self.material_classifier = nn.Linear(2048, num_material_classes)
+
+        # Dropout for regularization
+        self.dropout = nn.Dropout(0.3)
+
+    def forward(self, x):
+        # Extract features using backbone
+        features = self.backbone(x)
+        features = features.view(features.size(0), -1)
+        features = self.dropout(features)
+
+        # Get predictions for each attribute
+        object_pred = self.object_classifier(features)
+        material_pred = self.material_classifier(features)
 
-# Add the train directory to Python path
-train_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'train')
-sys.path.insert(0, train_dir)
+        return {
+            'object_name': object_pred,
+            'material': material_pred,
+        }
 
-# Now we can import from train.py
-try:
-    import train
-    MultiOutputModel = train.MultiOutputModel
-except ImportError as e:
-    print(f"Import error: {e}")
-    print("Make sure train.py exists in the train/ directory")
-    sys.exit(1)
+print("MultiOutputModel class defined directly in app (standalone)")
 
 class ArtifactClassifier:
     def __init__(self, model_path="train/outputs/best_model.pth"):
@@ -270,8 +297,9 @@ def create_interface():
         - Object type identification (coin, vase, statue, etc.)
         - Material classification (gold, silver, pottery, etc.)
         - Confidence scores for each prediction
-        - GPU-accelerated processing (RTX 2060)
+        - GPU-accelerated processing (if available)
         - Auto-downloads model from HuggingFace Hub
+        - Completely standalone - no training code needed
 
         Supported formats: JPG, PNG, JPEG
         """,