Update app.py

app.py

@@ -27,31 +27,36 @@ CLASS_RISK = {
 }
 COLORS = ['#2ecc71','#3498db','#e74c3c','#f39c12','#9b59b6']
 
-@spaces.GPU
+# ── Build model on CPU at startup ─────────────────────────────────────────────
 def build_model():
-    m = models.Sequential([
-        layers.Input(shape=(187, 1)),
-        layers.Conv1D(64, 7, padding='same', activation='relu', kernel_regularizer=regularizers.l2(1e-4)),
-        layers.BatchNormalization(), layers.MaxPooling1D(2), layers.Dropout(0.2),
-        layers.Conv1D(128, 5, padding='same', activation='relu', kernel_regularizer=regularizers.l2(1e-4)),
-        layers.BatchNormalization(), layers.MaxPooling1D(2), layers.Dropout(0.25),
-        layers.Conv1D(256, 3, padding='same', activation='relu', kernel_regularizer=regularizers.l2(1e-4)),
-        layers.BatchNormalization(), layers.MaxPooling1D(2), layers.Dropout(0.3),
-        layers.Conv1D(256, 3, padding='same', activation='relu', kernel_regularizer=regularizers.l2(1e-4)),
-        layers.BatchNormalization(), layers.GlobalAveragePooling1D(), layers.Dropout(0.3),
-        layers.Dense(256, activation='relu', kernel_regularizer=regularizers.l2(1e-4)),
-        layers.BatchNormalization(), layers.Dropout(0.4),
-        layers.Dense(128, activation='relu', kernel_regularizer=regularizers.l2(1e-4)),
-        layers.Dropout(0.3),
-        layers.Dense(5, activation='softmax')
-    ], name="CNN_ECG")
-    m.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
-    m.load_weights("cnn_weights.weights.h5")
+    with tf.device('/CPU:0'):
+        m = models.Sequential([
+            layers.Input(shape=(187, 1)),
+            layers.Conv1D(64, 7, padding='same', activation='relu',
+                          kernel_regularizer=regularizers.l2(1e-4)),
+            layers.BatchNormalization(), layers.MaxPooling1D(2), layers.Dropout(0.2),
+            layers.Conv1D(128, 5, padding='same', activation='relu',
+                          kernel_regularizer=regularizers.l2(1e-4)),
+            layers.BatchNormalization(), layers.MaxPooling1D(2), layers.Dropout(0.25),
+            layers.Conv1D(256, 3, padding='same', activation='relu',
+                          kernel_regularizer=regularizers.l2(1e-4)),
+            layers.BatchNormalization(), layers.MaxPooling1D(2), layers.Dropout(0.3),
+            layers.Conv1D(256, 3, padding='same', activation='relu',
+                          kernel_regularizer=regularizers.l2(1e-4)),
+            layers.BatchNormalization(), layers.GlobalAveragePooling1D(), layers.Dropout(0.3),
+            layers.Dense(256, activation='relu', kernel_regularizer=regularizers.l2(1e-4)),
+            layers.BatchNormalization(), layers.Dropout(0.4),
+            layers.Dense(128, activation='relu', kernel_regularizer=regularizers.l2(1e-4)),
+            layers.Dropout(0.3),
+            layers.Dense(5, activation='softmax')
+        ], name="CNN_ECG")
+        m.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
+        m.load_weights("cnn_weights.weights.h5")
     return m
 
 model = build_model()
 
-@spaces.GPU
+# ── Predict ───────────────────────────────────────────────────────────────────
 def predict_ecg(text_input):
     try:
         cleaned = text_input.replace(",", " ").replace("\n", " ").replace("\t", " ")
@@ -66,9 +71,12 @@ def predict_ecg(text_input):
 
         scaler        = MinMaxScaler()
         signal_scaled = scaler.fit_transform(features.reshape(-1, 1)).reshape(1, 187, 1)
-        probs         = model.predict(signal_scaled, verbose=0)[0]
-        pred_class    = int(np.argmax(probs))
-        confidence    = float(np.max(probs)) * 100
+
+        with tf.device('/CPU:0'):
+            probs = model.predict(signal_scaled, verbose=0)[0]
+
+        pred_class = int(np.argmax(probs))
+        confidence = float(np.max(probs)) * 100
 
         fig, axes = plt.subplots(1, 2, figsize=(14, 4))
         fig.patch.set_facecolor('#0e1117')
@@ -89,7 +97,8 @@ def predict_ecg(text_input):
 
         ax2 = axes[1]
         ax2.set_facecolor('#1a1a2e')
-        bars = ax2.barh([CLASS_MAPPING[i] for i in range(5)], probs*100, color=COLORS, alpha=0.85)
+        bars = ax2.barh([CLASS_MAPPING[i] for i in range(5)],
+                        probs * 100, color=COLORS, alpha=0.85)
         for bar, val in zip(bars, probs):
             ax2.text(val*100+0.5, bar.get_y()+bar.get_height()/2,
                      f"{val*100:.1f}%", va='center', color='white', fontsize=9)
@@ -102,10 +111,10 @@ def predict_ecg(text_input):
 
         plt.tight_layout()
 
-        result = f"**🤖 Predicted:** {CLASS_MAPPING[pred_class]}\n\n**📊 Confidence:** {confidence:.2f}%\n\n**⚕️ Risk:** {CLASS_RISK[pred_class]}"
+        result   = f"**🤖 Predicted:** {CLASS_MAPPING[pred_class]}\n\n**📊 Confidence:** {confidence:.2f}%\n\n**⚕️ Risk:** {CLASS_RISK[pred_class]}"
         true_out = ""
         if true_label is not None:
-            match = "✅ Correct!" if pred_class == true_label else "❌ Incorrect"
+            match    = "✅ Correct!" if pred_class == true_label else "❌ Incorrect"
             true_out = f"**🏷️ True Label:** {CLASS_MAPPING.get(true_label,'Unknown')} — {match}"
         stats = f"**Signal Stats:** R-peak={features.max():.4f} @ t={r_idx} | Min={features.min():.4f} | Mean={features.mean():.4f} | Std={features.std():.4f}"
 
@@ -114,6 +123,7 @@ def predict_ecg(text_input):
     except Exception as e:
         return None, f"❌ Error: {str(e)}", "", ""
 
+# ── UI ────────────────────────────────────────────────────────────────────────
 with gr.Blocks(title="ECG Classification") as demo:
     gr.Markdown("""
 # 🫀 ECG Arrhythmia Classification
@@ -122,8 +132,11 @@ with gr.Blocks(title="ECG Classification") as demo:
 """)
     with gr.Row():
         with gr.Column(scale=1):
-            text_input  = gr.Textbox(label="Paste 187 or 188 ECG values", lines=8,
-                                     placeholder="0.5, 0.8, 0.3, 1.0 ...")
+            text_input  = gr.Textbox(
+                label="Paste 187 or 188 ECG values",
+                lines=8,
+                placeholder="0.5, 0.8, 0.3, 1.0 ..."
+            )
             predict_btn = gr.Button("▶ Run Prediction", variant="primary")
             gr.Markdown("""
 | Label | Class | Risk |
@@ -140,9 +153,11 @@ with gr.Blocks(title="ECG Classification") as demo:
             true_out   = gr.Markdown()
             stats_out  = gr.Markdown()
 
-    predict_btn.click(predict_ecg, inputs=[text_input],
-                      outputs=[plot_out, result_out, true_out, stats_out])
-
+    predict_btn.click(
+        predict_ecg,
+        inputs=[text_input],
+        outputs=[plot_out, result_out, true_out, stats_out]
+    )
     gr.Markdown("*DEPI Final Project 2025*")
 
 demo.launch(ssr_mode=False)