Add second model

app.py

@@ -10,8 +10,12 @@ from PIL import Image
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 MODEL_PATH = os.path.join(BASE_DIR, "artifacts", "best_model.pt")
 LABELS_PATH = os.path.join(BASE_DIR, "data", "style_labels.json")
+# second model
+ARTIST_MODEL_PATH = os.path.join(BASE_DIR, "artifacts", "best_artist_model.pt")
+ARTIST_LABELS_PATH = os.path.join(BASE_DIR, "data", "artist_labels.json")
 
-class ArtStyleClassifier(nn.Module):
+# define the model architecture (same as training) to load the weights correctly, num_classes will be altered when loading
+class ArtClassifier(nn.Module): 
     def __init__(self, num_classes):
         super().__init__()
         self.backbone = timm.create_model("efficientnet_b0", pretrained=False, num_classes=0)
@@ -27,13 +31,18 @@ class ArtStyleClassifier(nn.Module):
         features = self.backbone(x)
         return self.classifier(features)
 
-with open(LABELS_PATH, "r") as f:
-    style_to_idx = json.load(f)
-idx_to_style = {int(i): s for s, i in style_to_idx.items()}
+# helper function to load model and labels, returns model and idx_to_label dict
+def load_model(model_path, labels_path):
+    with open(labels_path, "r") as f:
+        label_to_idx = json.load(f)
 
-model = ArtStyleClassifier(num_classes=len(style_to_idx))
-model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu")))
-model.eval()
+    idx_to_label = {int(i): label for label, i in label_to_idx.items()}
+
+    model = ArtClassifier(num_classes=len(label_to_idx))
+    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+    model.eval()
+
+    return model, idx_to_label
 
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -41,11 +50,19 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406],
                          std=[0.229, 0.224, 0.225]),
 ])
-def predict(image):
+
+# call the helper function to load both models and their corresponding idx_to_label mappings
+style_model, idx_to_style = load_model(MODEL_PATH, LABELS_PATH)
+artist_model, idx_to_artist = load_model(ARTIST_MODEL_PATH, ARTIST_LABELS_PATH)
+
+
+def make_prediction_html(model, idx_to_label, image, title):
     if image is None:
         return "<p>No image provided.</p>"
+    
     pil_image = Image.fromarray(image).convert("RGB")
     tensor = transform(pil_image).unsqueeze(0)
+
     with torch.no_grad():
         outputs = model(tensor)
         probs = torch.softmax(outputs, dim=1)[0]
@@ -53,14 +70,15 @@ def predict(image):
     
     medals = ["🥇", "🥈", "🥉"]
     colors = ["#4f46e5", "#7c3aed", "#a855f7"]
-    html = "<div style='font-family:sans-serif;padding:10px'>"
+    html = f"<div style='font-family:sans-serif;padding:10px'><h3>{title}</h3>"
+
     for i, (p, idx) in enumerate(zip(top_probs, top_idxs)):
-        style = idx_to_style[idx.item()]
+        label = idx_to_label[idx.item()] # rename variable to fit broad use case (style or artist)
         pct = float(p.item()) * 100
         html += f"""
         <div style='margin-bottom:16px'>
             <div style='display:flex;justify-content:space-between;margin-bottom:4px'>
-                <span style='font-weight:600'>{medals[i]} {style}</span>
+                <span style='font-weight:600'>{medals[i]} {label}</span>
                 <span style='color:#666'>{pct:.1f}%</span>
             </div>
             <div style='background:#e5e7eb;border-radius:999px;height:12px'>
@@ -70,12 +88,21 @@ def predict(image):
     html += "</div>"
     return html
 
+def predict(image):
+    style_html = make_prediction_html(style_model, idx_to_style, image, "Art Style Predictions")
+    artist_html = make_prediction_html(artist_model, idx_to_artist, image, "Artist Predictions")
+
+    return style_html, artist_html
+
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(label="Upload your artwork"),
-    outputs=gr.HTML(label="Art Style Predictions"),
-    title="Art Style Classifier",
-    description="Upload a photo of your artwork to identify which art style it most resembles.",
+    outputs=[
+        gr.HTML(label="Art Style Predictions"),
+        gr.HTML(label="Artist Predictions")
+    ],
+    title="Art Style + Artist Classifier",
+    description="Upload a photo of your artwork to identify which art style and artist's work it most resembles.",
 )
 
 demo.launch(server_name="0.0.0.0", server_port=7860)