Create app.py

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import torch
+import numpy as np
+
+# ЕҢ ТАЗА ЖӘНЕ ҚУАТТЫ МОДЕЛЬ ID-І (90%+ Accuracy)
+MODEL_ID = "keremberke/vit-base-patch16-224-full-empty-trash-bin" 
+CLASS_NAMES = ['Empty', 'Full'] 
+
+try:
+    # Модельді жүктеу кезінде қате болмауы үшін ең қарапайым әдісті қолданамыз
+    feature_extractor = ViTFeatureExtractor.from_pretrained(MODEL_ID)
+    model = ViTForImageClassification.from_pretrained(MODEL_ID)
+    MODEL_LOADED = True
+    if model.config.id2label:
+        CLASS_NAMES = [model.config.id2label[i] for i in model.config.id2label]
+
+except Exception as e:
+    MODEL_LOADED = False
+    
+def classify_trash_bin(image):
+    if not MODEL_LOADED:
+        # Енді "Модель жүктелмеді" дегеннің орнына "Error" деп шығады, бірақ жүктелуі керек
+        return {"Error": 1.0, "Check Logs": 0.0}
+    
+    if image is None:
+        return {CLASS_NAMES[0]: 0.5, CLASS_NAMES[1]: 0.5}
+    
+    try:
+        img = Image.fromarray(image).convert("RGB")
+        inputs = feature_extractor(images=img, return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+        
+        logits = outputs.logits
+        probabilities = torch.softmax(logits, dim=1).squeeze().tolist()
+        
+        if len(probabilities) > 2:
+             probabilities = probabilities[:2] 
+        
+        results = {CLASS_NAMES[i]: float(probabilities[i]) for i in range(len(CLASS_NAMES))}
+        return results
+
+    except Exception as e:
+        return {"Error": 1.0, "Check Logs": 0.0}
+
+# Gradio интерфейсін құру
+iface = gr.Interface(
+    fn=classify_trash_bin,
+    inputs=gr.Image(type="numpy", label="SmartTrachAI Input"),
+    outputs=gr.Label(num_top_classes=2, label="Prediction"),
+    title="SmartTrachAI", 
+    description="Automated Trash Bin Status Detector."
+)
+
+iface.launch()