chatbot/test.py

import os
import json
import torch
import sqlite3
import gradio as gr
import faiss
import pandas as pd
from datetime import datetime
from PIL import Image
from flask import Flask, jsonify, request
from werkzeug.utils import secure_filename
from threading import Thread
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, AutoTokenizer, AutoModel, \
    AutoModelForCausalLM
from matplotlib import pyplot as plt
from sklearn.linear_model import LinearRegression
import numpy as np
import seaborn as sns
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torch import nn, optim
from torch.optim import lr_scheduler
from sklearn.preprocessing import LabelEncoder


# Initialize Flask app and models
app = Flask(__name__, static_folder="static")

# Constants
LOW_STOCK_THRESHOLD = 5  # Customize threshold as needed
DATABASE = 'uploaded_images.db'

# Available model options
MODEL_OPTIONS = {
    "Google ViT (Base)": "google/vit-base-patch16-224",
    "Google ViT (Large)": "google/vit-large-patch16-224",
    "Microsoft ResNet50": "microsoft/resnet-50",
    "Facebook ConvNeXt Tiny": "facebook/convnext-tiny-224",
    "Microsoft Swin": "microsoft/swin-tiny-patch4-window7-224",
}

# Set default model
selected_model_name = MODEL_OPTIONS["Google ViT (Base)"]
feature_extractor = AutoFeatureExtractor.from_pretrained(selected_model_name)
model = AutoModelForImageClassification.from_pretrained(selected_model_name)
class_names = model.config.id2label

# Initialize inventory data
inventory_data = {}


# Initialize database
def init_db():
    conn = sqlite3.connect(DATABASE)
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS images (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT,
            upload_time TEXT
        )
    """)
    conn.commit()
    conn.close()


init_db()


def plot_bar_chart(inventory_data):
    # Get the list of items and their counts
    items = list(inventory_data.keys())
    counts = [inventory_data[item]["count"] for item in items]

    # Extract only the first word from each item name
    first_words = [item.split()[0] for item in items]

    # Create the bar chart
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.bar(first_words, counts, color="skyblue")

    # Set titles and labels
    ax.set_title("Inventory Counts")
    ax.set_xlabel("Items")
    ax.set_ylabel("Count")

    # Save the chart as a PNG file
    chart_path = "bar_chart.png"
    plt.tight_layout()
    plt.savefig(chart_path)
    plt.close()

    return chart_path



def plot_line_chart(inventory_data):
    # Extract items, counts, and dates from inventory data
    items = list(inventory_data.keys())
    counts = [inventory_data[item]["count"] for item in items]
    dates = [inventory_data[item]["last_detected"] for item in items]

    # Convert string dates to pandas datetime objects
    dates = pd.to_datetime(dates)

    # Create a DataFrame with dates and counts
    counts_df = pd.DataFrame({"Date": dates, "Count": counts})
    counts_df.sort_values("Date", inplace=True)

    # Format the 'Date' column into a human-readable format
    counts_df["Date"] = counts_df["Date"].dt.strftime('%Y-%m-%d %H:%M:%S')

    # Create the plot
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.plot(counts_df["Date"], counts_df["Count"], marker="o", color="orange")

    # Set the title and labels for the plot
    ax.set_title("Stock Changes Over Time")
    ax.set_xlabel("Date")
    ax.set_ylabel("Stock Level")

    # Save the plot as a PNG file
    chart_path = "line_chart.png"
    plt.xticks(rotation=45)  # Rotate x-ticks for better readability
    plt.tight_layout()
    plt.savefig(chart_path)
    plt.close()

    return chart_path


def plot_pie_chart(inventory_data):
    # Get the list of items and their counts
    items = list(inventory_data.keys())
    counts = [inventory_data[item]["count"] for item in items]

    # Extract only the first word from each item name
    first_words = [item.split()[0] for item in items]

    # Create the pie chart
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.pie(counts, labels=first_words, autopct='%1.1f%%', startangle=140, colors=plt.cm.Paired.colors)

    # Set the title
    ax.set_title("Product Category Breakdown")

    # Save the chart as a PNG file
    chart_path = "pie_chart.png"
    plt.tight_layout()
    plt.savefig(chart_path)
    plt.close()

    return chart_path



def plot_heatmap(inventory_data):
    # Get the list of items
    items = list(inventory_data.keys())

    # Extract only the first word from each item name
    first_words = [item.split()[0] for item in items]

    # Create a matrix where the count is placed at the correct location
    change_matrix = [[inventory_data[item]["count"] if item == other else 0 for other in items] for item in items]

    # Create the heatmap
    fig, ax = plt.subplots(figsize=(8, 6))
    cax = ax.imshow(change_matrix, cmap="YlOrRd", interpolation='nearest')
    # Set the labels on the x and y axes to the first words of the items
    ax.set_xticks(range(len(items)))
    ax.set_yticks(range(len(items)))
    ax.set_xticklabels(first_words)
    ax.set_yticklabels(first_words)
    # Add a colorbar
    fig.colorbar(cax)
    # Set the title of the heatmap
    ax.set_title("Stock Change Heatmap")
    # Save the chart as a PNG file
    chart_path = "heatmap.png"
    plt.tight_layout()
    plt.savefig(chart_path)
    plt.close()

    return chart_path



# Function to return the paths of the images
def generate_bar_chart():
    return plot_bar_chart(inventory_data)


def generate_line_chart():
    return plot_line_chart(inventory_data)


def generate_pie_chart():
    return plot_pie_chart(inventory_data)


def generate_heatmap():
    return plot_heatmap(inventory_data)


# Utility functions
def log_inventory_data(item_class):
    try:
        timestamp = datetime.now().isoformat()

        # Initialize the inventory data for the item class if it's not present
        if item_class not in inventory_data:
            inventory_data[item_class] = {
                "category": item_class,
                "count": 0,  # Initial count
                "last_detected": None,  # Timestamp of last detection
                "history": []  # Track the historical counts over time
            }

        # Log the count and timestamp to the history
        inventory_data[item_class]["history"].append({
            "timestamp": timestamp,  # Store the timestamp as a string
            "count": inventory_data[item_class]["count"]
        })

        # Update the count of the item class
        inventory_data[item_class]["count"] += 1

        # Update the last detected timestamp
        inventory_data[item_class]["last_detected"] = timestamp

        # Optionally: Save the data to a file for persistence
        with open("inventory_log.json", "w") as f:
            json.dump(inventory_data, f, indent=4)

        print(f"Inventory data logged for item class: {item_class}")

    except Exception as e:
        print(f"Error logging inventory data: {e}")


def check_stock_levels():
    try:
        for item, details in inventory_data.items():
            if details["count"] < LOW_STOCK_THRESHOLD:
                print(f"Low stock detected for {item}: {details['count']} items.")
    except Exception as e:
        print(f"Error checking stock levels: {e}")


def save_image_to_db(file):
    filename = secure_filename(file.filename)
    upload_time = datetime.now().isoformat()
    file_path = os.path.join("uploads", filename)
    file.save(file_path)
    conn = sqlite3.connect(DATABASE)
    cursor = conn.cursor()
    cursor.execute("INSERT INTO images (filename, upload_time) VALUES (?, ?)", (filename, upload_time))
    conn.commit()
    conn.close()


def batch_predict(images):
    results = []
    try:
        for image_file in images:
            with Image.open(image_file.name) as image:
                if image.mode != "RGB":
                    image = image.convert("RGB")
                inputs = feature_extractor(images=image, return_tensors="pt")
                outputs = model(**inputs)
                predicted_class_id = torch.argmax(outputs.logits, dim=1).item()
                item_class = class_names[predicted_class_id]
                log_inventory_data(item_class)  # Pass only the item_class argument
                results.append({"Image": os.path.basename(image_file.name), "Classification": item_class})
    except Exception as e:
        return [{"Image": "Error", "Classification": str(e)}]
    return pd.DataFrame(results)


def forecast_inventory(item_class, days=7):
    """Use linear regression for basic inventory forecasting."""
    try:
        # Get the historical data for the item class
        if item_class not in inventory_data or not inventory_data[item_class]["history"]:
            return {"error": f"No historical data for {item_class}."}

        history = inventory_data[item_class]["history"]
        timestamps = [datetime.fromisoformat(entry["timestamp"]) for entry in history]
        counts = [entry["count"] for entry in history]

        # Convert timestamps to days since the first entry
        days_since_first_entry = [(timestamp - timestamps[0]).days for timestamp in timestamps]

        # Apply linear regression to forecast the next `days` values
        model = LinearRegression()
        model.fit(np.array(days_since_first_entry).reshape(-1, 1), counts)

        # Predict future inventory counts
        future_days = np.array([days_since_first_entry[-1] + i for i in range(1, days + 1)]).reshape(-1, 1)
        predictions = model.predict(future_days)

        forecast = [{"day": i + 1, "predicted_count": int(predictions[i])} for i in range(days)]
        return forecast
        print(f"Inventory data logged for item class: {forecast}")

    except Exception as e:
        return {"error": f"Error in forecasting: {str(e)}"}


def change_model(selected_model_key):
    try:
        global feature_extractor, model, class_names

        if selected_model_key not in MODEL_OPTIONS:
            return "Error: Invalid model selection"

        # Retrieve the model path (either pre-trained or fine-tuned)
        selected_model_path = MODEL_OPTIONS[selected_model_key]

        # Load the model from the specified path
        feature_extractor = AutoFeatureExtractor.from_pretrained(selected_model_path)
        model = AutoModelForImageClassification.from_pretrained(selected_model_path)
        class_names = model.config.id2label

        return f"Model changed to {selected_model_key}"
    except Exception as e:
        return f"Error changing model: {str(e)}"


# Define a custom dataset class
class CustomDataset(Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = Image.open(self.images[idx]).convert("RGB")
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

# Define transformations for image preprocessing
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Function to preprocess and load data
def preprocess_data(train_images, train_labels):
    dataset = CustomDataset(train_images, train_labels, transform=transform)
    dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
    return dataloader


# Fine-tune the model with a custom dataset
def fine_tune_model(train_images, train_labels, model_name="custom_model"):
    try:
        # Convert image paths to a suitable format
        train_images = [image.name for image in train_images]  # list of image paths

        # Check if labels need encoding
        if isinstance(train_labels, str):
            train_labels = train_labels.split(",")  # Convert the comma-separated string to a list

        # Encode labels if they are not integers
        label_encoder = LabelEncoder()
        train_labels = label_encoder.fit_transform(train_labels)

        # Load the model and prepare for fine-tuning
        global model
        model = AutoModelForImageClassification.from_pretrained(selected_model_name)

        # Modify the classifier layer for the new dataset
        num_labels = len(set(train_labels))  # Number of unique labels
        model.classifier = nn.Linear(model.config.hidden_size, num_labels)

        # Prepare data loader
        dataloader = preprocess_data(train_images, train_labels)

        # Define loss function, optimizer, and learning rate scheduler
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=0.0001)
        scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

        # Training loop
        num_epochs = 50  # Set the number of epochs
        model.train()  # Set model to training mode
        for epoch in range(num_epochs):
            running_loss = 0.0
            for inputs, labels in dataloader:
                optimizer.zero_grad()

                # Forward pass
                outputs = model(inputs)
                logits = outputs.logits if hasattr(outputs, "logits") else outputs  # Handle outputs

                # Calculate loss
                loss = criterion(logits, labels)
                loss.backward()
                optimizer.step()

                running_loss += loss.item()

            scheduler.step()
            print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(dataloader)}")
            # return f"Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(dataloader)}"

        # Save the fine-tuned model
        fine_tuned_model_path = f"models/{model_name}"  # Save path
        model.save_pretrained(fine_tuned_model_path)
        print(f"Fine-tuned model saved to {fine_tuned_model_path}")

        # Dynamically add the fine-tuned model to the model selection
        MODEL_OPTIONS[model_name] = fine_tuned_model_path
        return f"Fine-tuned model {model_name} has been added successfully!"
    except Exception as e:
        return f"Error fine-tuning the model: {str(e)}"



proxy_prefix = os.environ.get("PROXY_PREFIX")

# Gradio Interface
with gr.Blocks() as interface:
    gr.Markdown("## VisionTrack - Smart Inventory Management and Analysis System")

    with gr.Tab("Model Selection/Fine-Tuning"):
        with gr.Tab("Select Model"):
            gr.Markdown("Choose a model for image classification.")
            model_dropdown = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Select Model",
                                         value="Google ViT (Base)")
            model_dropdown.change(change_model, inputs=model_dropdown, outputs=gr.Textbox(label="status"))
        with gr.Tab("Model Fine-Tuning"):
            gr.Markdown("Upload your custom dataset for model fine-tuning.")
            train_image_input = gr.File(label="Upload Training Images", file_count="multiple", type="filepath")
            train_label_input = gr.Textbox(label="Enter Corresponding Labels (comma-separated)")
            fine_tune_button = gr.Button("Fine-Tune Model")
            fine_tune_output = gr.Text(label="Training Status")
            fine_tune_button.click(fine_tune_model, inputs=[train_image_input, train_label_input],
                                   outputs=fine_tune_output)

    with gr.Tab("Image Classification"):
        gr.Markdown("Upload images for classification.")
        image_input = gr.File(label="Upload Images", file_count="multiple", type="filepath")
        output = gr.DataFrame(label="Classification Results", interactive=True)
        image_input.upload(batch_predict, inputs=image_input, outputs=output)

    with gr.Tab("Inventory"):
        gr.Markdown("Check current inventory.")
        inventory_display = gr.DataFrame(value=pd.DataFrame(columns=["Item", "Count", "Last Detected"]))
        gr.Button("Fetch Inventory").click(lambda: pd.DataFrame(inventory_data).T, outputs=inventory_display)

    with gr.Tab("Forecasting"):
        gr.Markdown("Forecast inventory levels for the next 7 days.")
        item_class_input = gr.Textbox(label="Enter Item Class")
        forecast_button = gr.Button("Forecast")
        forecast_output = gr.Text(label="Predicted Inventory Forecast")
        forecast_button.click(forecast_inventory, inputs=item_class_input, outputs=forecast_output)

    with gr.Tab("Inventory Dashboard"):
        with gr.Tab("Bar Chart"):
            bar_chart_image = gr.Image(type="filepath", label="Bar Chart")
            gr.Button("Generate Bar Chart").click(generate_bar_chart, outputs=bar_chart_image)

        with gr.Tab("Line Chart"):
            line_chart_image = gr.Image(type="filepath", label="Line Chart")
            gr.Button("Generate Line Chart").click(generate_line_chart, outputs=line_chart_image)

        with gr.Tab("Pie Chart"):
            pie_chart_image = gr.Image(type="filepath", label="Pie Chart")
            gr.Button("Generate Pie Chart").click(generate_pie_chart, outputs=pie_chart_image)

        with gr.Tab("Heatmap"):
            heatmap_image = gr.Image(type="filepath", label="Heatmap")
            gr.Button("Generate Heatmap").click(generate_heatmap, outputs=heatmap_image)

if __name__ == '__main__':
    app_thread = Thread(target=lambda: app.run(debug=True, use_reloader=False))
    app_thread.start()
    interface.launch(server_name="0.0.0.0", server_port=8080, root_path=proxy_prefix, share=True)