Upload inference.py via DNA Console (Portable Version)

inference.py

@@ -125,7 +125,15 @@ class SequenceFeatureExtractor:
             for b in 'ATGC': res.append(p_c[p].get(b, 0) / t)
         return res
 
-def predict_dna(sequence):
+def predict_dna(sequence, confidence_threshold=0.55, rare_class_threshold=0.65):
+    """
+    DNA sequence prediction with confidence thresholds.
+    
+    Args:
+        sequence: DNA sequence string
+        confidence_threshold: Minimum confidence for general classification (default 55%)
+        rare_class_threshold: Higher threshold for rare classes like Influenza B (default 65%)
+    """
     # Load Models
     rf_model = joblib.load("dna_classifier.joblib")
     rf_scaler = joblib.load("scaler_rf.joblib")
@@ -137,16 +145,41 @@ def predict_dna(sequence):
     feat_rf = extractor_rf.transform([sequence])
     scaled_rf = rf_scaler.transform(feat_rf)
     type_basic = rf_model.predict(scaled_rf)[0]
+    rf_proba = rf_model.predict_proba(scaled_rf)[0]
+    rf_confidence = max(rf_proba)
     
-    # 2. ViralBoost Prediction (Virus Type)
+    # 2. ViralBoost Prediction (Virus Type) with Confidence Check
     extractor_gb = SequenceFeatureExtractor()
     feat_gb = extractor_gb.transform([sequence])
     scaled_gb = gb_scaler.transform(feat_gb)
-    type_virus = gb_model.predict(scaled_gb)[0]
+    
+    gb_proba = gb_model.predict_proba(scaled_gb)[0]
+    gb_confidence = max(gb_proba)
+    predicted_idx = gb_proba.argmax()
+    predicted_class = gb_model.classes_[predicted_idx]
+    
+    # 희귀 클래스 (Influenza B 등)는 더 높은 신뢰도 요구
+    rare_classes = ['Influenza B', 'Chicken anemia virus']
+    if predicted_class in rare_classes:
+        effective_threshold = rare_class_threshold
+    else:
+        effective_threshold = confidence_threshold
+    
+    # 신뢰도 임계값 미달 시 'Unknown'으로 분류
+    if gb_confidence < effective_threshold:
+        type_virus = 'Unknown'
+        virus_confidence = gb_confidence
+    else:
+        type_virus = predicted_class
+        virus_confidence = gb_confidence
     
     return {
         "classification": type_basic,
-        "virus_identity": type_virus
+        "classification_confidence": float(rf_confidence),
+        "virus_identity": type_virus,
+        "virus_confidence": float(virus_confidence),
+        "raw_prediction": predicted_class,  # 원래 예측 (디버깅용)
+        "raw_confidence": float(gb_confidence)
     }
 
 if __name__ == "__main__":