Labels Added

app.py

@@ -2,20 +2,29 @@ import streamlit as st
 import torch
 from PIL import Image
 from torchvision import transforms
+import json
 from src.model import SwinTransformerMultiLabel  # Import from src folder
 
 # Title and description
 st.title("STAR Multi-Label Classifier")
-st.write("Upload an image to get predictions.")
+st.write("Upload an image to get multi-label predictions.")
+
+# Load class labels from JSON
+label_file = "data/labels.json"  # Path to labels file
+with open(label_file, "r") as f:
+    label_data = json.load(f)
+
+# Extract unique class labels
+class_labels = sorted(set(tag for tags in label_data.values() for tag in tags))
+NUM_CLASSES = len(class_labels)
 
 # Load trained model
-NUM_CLASSES = 18
-model_path = "models/star.pth"  
+model_path = "models/star.pth"
 model = SwinTransformerMultiLabel(num_classes=NUM_CLASSES)
 model.load_state_dict(torch.load(model_path, map_location="cpu"))
 model.eval()
 
-# Define transformation
+# Define image preprocessing transformations
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -32,7 +41,8 @@ if uploaded_file is not None:
     img_tensor = transform(image).unsqueeze(0)
     with torch.no_grad():
         output = model(img_tensor)
-        predicted_labels = (output > 0.5).int().tolist()[0]
+        predicted_indices = [i for i in range(NUM_CLASSES) if output[0][i] > 0.5]  # Threshold = 0.5
+        predicted_labels = [class_labels[i] for i in predicted_indices]  # Convert indices to labels
 
-    # Display Results
-    st.write("✅ **Predicted Labels:**", predicted_labels)
+    # Display results
+    st.write("✅ **Predicted Labels:**", ", ".join(predicted_labels) if predicted_labels else "No labels detected")