Integrate FastAI model for warning lamp detection and update dependencies

.gitignore

@@ -0,0 +1,20 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Jupyter Notebook
+.ipynb_checkpoints
+Inference_notebook_demo.ipynb
+
+# Virtual Environment
+venv/
+env/
+.env/
+
+# IDE
+.vscode/
+.idea/
+
+# Misc
+.DS_Store 

app.py

@@ -1,40 +1,56 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from fastai.vision.all import *
+from fastai.learner import load_learner
+from pathlib import Path
 import os
 
 """
-Warning Lamp Detector using Hugging Face Inference API
+Warning Lamp Detector using FastAI
 This application allows users to upload images of warning lamps and get classification results.
 """
 
-# Initialize the client with your model
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# Load the FastAI model
+try:
+    model_path = Path("WarningLampClassifier.pkl")
+    learn_inf = load_learner(model_path)
+    print("Model loaded successfully")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
 
 def detect_warning_lamp(image, history: list[tuple[str, str]], system_message):
     """
-    Process the uploaded image and return detection results
+    Process the uploaded image and return detection results using FastAI model
+    Args:
+        image: PIL Image from Gradio
+        history: Chat history
+        system_message: System prompt
+    Returns:
+        Updated chat history with prediction results
     """
-    # TODO: Replace with actual model inference
-    # This is a placeholder response - you'll need to integrate your actual model
-    messages = [{"role": "system", "content": system_message}]
-    
-    # Add the image analysis request
-    messages.append({
-        "role": "user", 
-        "content": f"Please analyze this warning lamp image and provide a detailed classification."
-    })
-
-    response = ""
-    for message in client.chat_completion(
-        messages,
-        max_tokens=512,
-        stream=True,
-        temperature=0.7,
-        top_p=0.95,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        yield response
+    try:
+        # Convert PIL image to FastAI compatible format
+        img = PILImage(image)
+        
+        # Get model prediction
+        pred_class, pred_idx, probs = learn_inf.predict(img)
+        
+        # Format the prediction results
+        confidence = float(probs[pred_idx])  # Convert to float for better formatting
+        response = f"Detected Warning Lamp: {pred_class}\nConfidence: {confidence:.2%}"
+        
+        # Add probabilities for all classes
+        response += "\n\nProbabilities for all classes:"
+        for i, (cls, prob) in enumerate(zip(learn_inf.dls.vocab, probs)):
+            response += f"\n- {cls}: {float(prob):.2%}"
+            
+        # Update chat history
+        history.append((None, response))
+        return history
+    except Exception as e:
+        error_msg = f"Error processing image: {str(e)}"
+        history.append((None, error_msg))
+        return history
 
 # Create a custom interface with image upload
 with gr.Blocks(title="Warning Lamp Detector", theme=gr.themes.Soft()) as demo:
@@ -46,8 +62,14 @@ with gr.Blocks(title="Warning Lamp Detector", theme=gr.themes.Soft()) as demo:
     1. Upload a clear image of the warning lamp
     2. Wait for the analysis
     3. View the detailed classification results
+    
+    ### Supported Warning Lamps:
     """)
     
+    # Display supported classes if available
+    if 'learn_inf' in locals():
+        gr.Markdown("\n".join([f"- {cls}" for cls in learn_inf.dls.vocab]))
+    
     with gr.Row():
         with gr.Column(scale=1):
             image_input = gr.Image(
@@ -58,7 +80,8 @@ with gr.Blocks(title="Warning Lamp Detector", theme=gr.themes.Soft()) as demo:
             system_message = gr.Textbox(
                 value="You are an expert in warning lamp classification. Analyze the image and provide detailed information about the type, color, and status of the warning lamp.",
                 label="System Message",
-                lines=3
+                lines=3,
+                visible=False  # Hide this since we're using direct model inference
             )
         
         with gr.Column(scale=1):

requirements.txt

@@ -1,3 +1,6 @@
 gradio>=4.19.2
 huggingface-hub>=0.20.3
-Pillow>=10.0.0  # Required for image processing
+Pillow>=10.0.0  # Required for image processing
+fastai>=2.7.13  # Required for model inference
+torch>=2.2.0    # Required by FastAI
+torchvision>=0.17.0  # Required by FastAI