added necessary files for app to run on HF

app.py

@@ -0,0 +1,256 @@
+# --- START OF FILE app_main.py ---
+
+import streamlit as st
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+# Ensure classification_model is importable (e.g., it's in the same directory orPYTHONPATH)
+try:
+    from classification_model.model import VGG16Model
+except ImportError:
+    st.error("Could not import VGG16Model. Make sure 'classification_model/model.py' exists and is importable.")
+    st.stop()
+
+from PIL import Image
+import os
+from dotenv import load_dotenv
+from langchain.schema import SystemMessage, HumanMessage, AIMessage, BaseMessage
+import io
+
+# Import graph logic from the other file
+try:
+    from graph_logic import initialize_llm, create_chat_graph, GraphState
+except ImportError:
+     st.error("Could not import from graph_logic.py. Make sure the file exists in the same directory.")
+     st.stop()
+
+
+# Load environment variables from .env file
+load_dotenv()
+
+# --- Helper Functions (Specific to App/UI) ---
+
+@st.cache_resource
+def load_vgg16(path='./classification_model/best_model.pth', device='cpu'):
+    """Loads the pre-trained VGG16 model."""
+    model = VGG16Model()
+    if not os.path.exists(path):
+        st.error(f"Model file not found at {path}. Please ensure the model is present.")
+        st.stop()
+    try:
+        model.load_state_dict(torch.load(path, map_location=device))
+        model.eval()
+    except Exception as e:
+        st.error(f"Error loading model state dict: {e}")
+        st.stop()
+    return model
+
+@st.cache_resource
+def load_class_labels() -> dict:
+    """Loads the class labels."""
+    return {0: 'good_images', 1: 'Imaging Artifact', 2: 'Not Tracking', 3: 'Tip Contamination'}
+
+def preprocess_image(img: Image.Image):
+    """Preprocesses the image for the VGG16 model."""
+    preprocess = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.3718, 0.1738, 0.0571], std=[0.2095, 0.2124, 0.1321]),
+    ])
+    img_tensor = preprocess(img).unsqueeze(0)
+    return img_tensor
+
+def predict_image_class(img: Image.Image, model, class_names: dict) -> str:
+    """Predicts the class label for a given image."""
+    img_tensor = preprocess_image(img)
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = F.softmax(outputs, dim=1)
+        top_prob, top_idx = torch.topk(probs, 1)
+        class_label = class_names.get(top_idx.item(), "Unknown Class")
+    return class_label
+
+def download_chat_history() -> str:
+    """Generates the chat history text for download."""
+    if "messages" not in st.session_state or not st.session_state.messages:
+        return ""
+    output = io.StringIO()
+    start_index = 1 if isinstance(st.session_state.messages[0], SystemMessage) else 0
+    for msg in st.session_state.messages[start_index:]:
+        role = "User" if isinstance(msg, HumanMessage) else "Assistant"
+        output.write(f"{role}: {msg.content}\n")
+    return output.getvalue()
+
+# --- Streamlit UI ---
+
+st.set_page_config(page_title="AFM Defect Assistant (LangGraph)", page_icon="🔬")
+
+st.title("🔬 AFM Image Defect Classification + LLM-based AFM Assistant")
+st.write("Upload an AFM image, get a classification, and chat with an AI assistant about the result.")
+
+# --- Sidebar Controls ---
+st.sidebar.header("Settings")
+
+# Model Selection
+provider = st.sidebar.selectbox("LLM Provider", ["OpenAI", "Anthropic"], key="provider_select")
+
+api_key = None
+api_key_name = ""
+if provider == "OpenAI":
+    default_model = "gpt-4o"
+    available_models = ["gpt-4o", "o3-mini"]
+    api_key = os.getenv("OPENAI_API_KEY")
+    api_key_name = "OPENAI_API_KEY"
+elif provider == "Anthropic":
+    default_model = "claude-3-5-sonnet-latest"
+    available_models = ["claude-3-5-sonnet-latest", "claude-3-7-sonnet-latest"]
+    api_key = os.getenv("ANTHROPIC_API_KEY")
+    api_key_name = "ANTHROPIC_API_KEY"
+else:
+    st.sidebar.error("Invalid provider selected.")
+    st.stop()
+
+# Display warning if API key is missing
+if not api_key:
+    st.sidebar.warning(f"{provider} API key not found. Please set the {api_key_name} environment variable.")
+
+model_name = st.sidebar.selectbox(f"Choose {provider} Model", available_models, index=available_models.index(default_model), key="model_name_select")
+temperature = st.sidebar.slider("LLM Temperature", 0.0, 1.0, 0.3, 0.05, key="temp_slider")
+
+# Clear Chat Button
+if st.sidebar.button("Start New Session", key="clear_chat_button"):
+    keys_to_clear = ["messages", "current_label", "uploaded_file_state", "llm", "graph"]
+    for key in keys_to_clear:
+        if key in st.session_state:
+            del st.session_state[key]
+    st.rerun()
+
+# --- Main Page Logic ---
+
+# File Uploader
+uploaded_file = st.file_uploader("Upload an AFM image", type=["jpg", "jpeg", "png"], key="file_uploader")
+
+# Manage state based on uploaded file
+if uploaded_file is not None:
+    new_file_id = uploaded_file.file_id
+    # Check if it's a new file or the same one to avoid re-processing
+    if "uploaded_file_state" not in st.session_state or st.session_state.uploaded_file_state["id"] != new_file_id:
+        st.session_state.uploaded_file_state = {"id": new_file_id, "name": uploaded_file.name}
+        # Clear previous chat/state if a new file is uploaded
+        keys_to_reset = ["messages", "current_label", "llm", "graph"]
+        for key in keys_to_reset:
+            if key in st.session_state:
+                del st.session_state[key]
+
+    # --- Image Processing and Classification ---
+    try:
+        img = Image.open(uploaded_file).convert("RGB")
+        st.image(img, caption=f"Uploaded: {st.session_state.uploaded_file_state['name']}", width=200)
+
+        model = load_vgg16()
+        class_names = load_class_labels()
+
+        with st.spinner("Classifying image..."):
+            class_label = predict_image_class(img, model, class_names)
+        st.success(f"**Predicted Class Label:** {class_label}")
+
+        # --- LLM and Graph Initialization ---
+        label_changed = ("current_label" not in st.session_state or
+                         st.session_state.current_label != class_label)
+
+        # Initialize LLM and Graph if not present or if label changed
+        if "llm" not in st.session_state or "graph" not in st.session_state or label_changed:
+            if not api_key:
+                st.error(f"Cannot proceed without {api_key_name}. Please set it in your environment variables.")
+                st.stop()
+            try:
+                # Initialize LLM using the function from graph_logic.py
+                st.session_state.llm = initialize_llm(provider, model_name, temperature, api_key)
+                # Create the graph using the function from graph_logic.py
+                st.session_state.graph = create_chat_graph(st.session_state.llm)
+                st.session_state.current_label = class_label
+
+                # Define the system prompt and initial state for the graph
+                system_prompt_content = (
+                    f"You are an expert in atomic force microscopy (AFM). "
+                    f"The user has uploaded an image and it has '{class_label}' defect. "
+                    "Your role is to help the user understand this defect, potential causes, "
+                    "and how to potentially avoid or address the issue represented by this defect. "
+                    "Provide concise, technically accurate, and helpful answers. Avoid speculation if unsure."
+                )
+                system_message = SystemMessage(content=system_prompt_content)
+                st.session_state.messages = [system_message] # Initialize message history
+
+            except Exception as e:
+                st.error(f"Failed to initialize LLM or Graph: {e}")
+                st.stop()
+
+        # --- Chat Interface ---
+        st.divider()
+        st.header(f"Chat about '{st.session_state.current_label}'")
+
+        # Display existing messages (skip system message)
+        if "messages" in st.session_state:
+            for i, msg in enumerate(st.session_state.messages):
+                if i == 0 and isinstance(msg, SystemMessage):
+                    continue
+                with st.chat_message("user" if isinstance(msg, HumanMessage) else "assistant"):
+                    st.markdown(msg.content)
+
+        # Chat input
+        if prompt := st.chat_input("Ask a question about the detected defect..."):
+            # Add user message to state and display it
+            st.session_state.messages.append(HumanMessage(content=prompt))
+            with st.chat_message("user"):
+                st.markdown(prompt)
+
+            # Prepare the input state for the graph
+            current_graph_state: GraphState = {"messages": st.session_state.messages}
+
+            # Invoke the graph
+            with st.chat_message("assistant"):
+                with st.spinner("Thinking..."):
+                    try:
+                        # Invoke the graph with the current state
+                        response_state = st.session_state.graph.invoke(current_graph_state)
+
+                        # Update session state with the full response history from the graph
+                        st.session_state.messages = response_state['messages']
+                        ai_response_content = st.session_state.messages[-1].content
+                        st.markdown(ai_response_content)
+
+                    except Exception as e:
+                        st.error(f"Error during chat generation: {e}")
+                        # Roll back user message if AI fails
+                        if st.session_state.messages and isinstance(st.session_state.messages[-1], HumanMessage):
+                            st.session_state.messages.pop()
+
+
+        # --- Download Chat Button ---
+        if len(st.session_state.get("messages", [])) > 1: # Show only if conversation started
+             st.divider()
+             chat_text = download_chat_history()
+             st.download_button(
+                 label="Download Chat History",
+                 data=chat_text,
+                 file_name=f"afm_chat_{st.session_state.current_label.replace(' ', '_')}.txt",
+                 mime="text/plain"
+             )
+
+    except Exception as e:
+        st.error(f"An error occurred processing the image or during chat setup: {e}")
+        if "uploaded_file_state" in st.session_state:
+            del st.session_state.uploaded_file_state # Reset if critical error occurs
+
+elif "uploaded_file_state" in st.session_state:
+     # If file uploader is cleared by the user after a file was processed
+     keys_to_clear = ["messages", "current_label", "uploaded_file_state", "llm", "graph"]
+     for key in keys_to_clear:
+         if key in st.session_state:
+             del st.session_state[key]
+     st.rerun()
+
+else:
+    st.info("Please upload an image to start the analysis and chat.")
+
+# --- END OF FILE app_main.py ---

classification_model/model.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+from torchvision.models import VGG16_Weights
+
+class VGG16Model(nn.Module):
+    def __init__(self, num_classes=4):
+       super(VGG16Model, self).__init__()
+       #self.vgg16 = models.vgg16(pretrained=True)
+       # Replace pretrained=True with weights=VGG16_Weights.DEFAULT
+       #self.vgg16 = models.vgg16(weights=VGG16_Weights.DEFAULT)
+       # Manually load the weights
+       self.vgg16 = models.vgg16()
+    #    state_dict = torch.load("/work/mech-ai/angona3/Trial/vgg16_weights.pth", map_location=torch.device("cpu"))
+    #    self.vgg16.load_state_dict(state_dict)
+       self.vgg16.classifier[6] = nn.Linear(4096, num_classes) ##In PyTorch CrossEntropyLoss Handles Softmax Internally
+       ###Adding an explicit softmax layer would create numerical instability by double-applying softmax (once in your model, once in the loss), 
+       # leading to incorrect gradient calculations during training.
+
+       self.vgg16.classifier[2] = nn.Dropout(p=0.6)  # or 0.7 ##introducing dropout for regularization
+       self.vgg16.classifier[5] = nn.Dropout(p=0.6)  # or 0.7
+
+       # Freeze all the convolutional layers (feature extractor part)
+       # The classifier layers (fully connected layers) remain trainable
+       for param in self.vgg16.features.parameters():
+           param.requires_grad = False
+           #param.requires_grad = True  # Unfreeze the last two fully connected layers
+        
+         # Unfreeze Conv Block 4 and Conv Block 5 (512 filters, 3x3 filters, same padding) 
+        # best: conv_layers_to_unfreeze = [17, 19, 21, 24, 26, 28] 
+       conv_layers_to_unfreeze = [17, 19, 21, 24, 26, 28] 
+       for layer_idx in conv_layers_to_unfreeze: 
+          for param in self.vgg16.features[layer_idx].parameters(): 
+               param.requires_grad = True
+
+
+
+
+       # Unfreeze the last two fully connected layers
+       # Unfreeze all fully connected layers
+       for param in self.vgg16.classifier.parameters():
+            param.requires_grad = True
+
+    def forward(self, x):
+        return self.vgg16(x)
+
+#Fine-tuning the entire network can lead to better performance compared to freezing layers 
+##because the model can adjust both the feature extractor and classifier to your specific dataset.
+##Yes, nn.CrossEntropyLoss in PyTorch is explicitly designed to handle raw logits (ranging from −∞to +∞) directly and efficiently. 
+##CrossEntropyLoss=Softmax(logits)+Log+NLLLoss
+
+##### In newer versions of torchvision, the weights argument replaces pretrained. 
+# #The VGG16_Weights.DEFAULT is the new way to specify that anyone want the pretrained weights, and it is the preferred method.
+##Deprecation Warning Fix: Replace pretrained=True with weights=VGG16_Weights.DEFAULT. 
+##Security Warning Fix: You don't need to use torch.load in this case because you're not loading a pre-trained model from a .pth file.
+##from torchvision.models import VGG16_Weights
+## self.vgg16 = models.vgg16(weights=VGG16_Weights.DEFAULT)
+
+
+## adding dropout to VGG16
+#class VGG16Model(nn.Module):
+#    def __init__(self, num_classes=4):
+#        super(VGG16Model, self).__init__()
+#        self.vgg16 = models.vgg16(pretrained=True)
+#        self.vgg16.classifier[6] = nn.Sequential(
+#            nn.Linear(4096, 1024),  # First layer reduces dimensions to 1024
+#            nn.ReLU(),             # Adds non-linearity to increase learning capacity
+#            nn.Dropout(0.5),       # Introduces dropout for regularization
+#            nn.Linear(1024, num_classes)  # Second layer maps to the desired number of classes
+#        )
+
+#    def forward(self, x):
+#        return self.vgg16(x)
+
+
+#import torch.nn as nn
+#from torchvision import models
+
+#def get_resnet18_model(num_classes):
+#    model = models.resnet18(pretrained=True)
+#    model.fc = nn.Linear(model.fc.in_features, num_classes)
+#    return model
+
+##VGG16 structure
+#(vgg16.classifier): Sequential(
+#  (0): Linear(25088, 4096)
+#  (1): ReLU(inplace=True)
+#  (2): Dropout(p=0.5, inplace=False)
+#  (3): Linear(4096, 4096)
+#  (4): ReLU(inplace=True)
+#  (5): Dropout(p=0.5, inplace=False)
+#  (6): Linear(4096, num_classes)
+#)

classification_model/predict.py

@@ -0,0 +1,52 @@
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+from model import VGG16Model
+# Load your PyTorch VGG16 model (pre-trained)
+
+def load_vgg16():
+    model = VGG16Model()
+    model.load_state_dict(torch.load("/work/mech-ai/jrrade/AFM/AFM-LLM-Defect-Guidance/classification_model/best_model.pth", map_location=torch.device('cpu')))  # Or GPU if available
+    model.eval()
+    return model
+
+# Load ImageNet class labels (you need a mapping file)
+
+def load_class_labels():
+    class_names = {0:'good_images', 1:'Imaging Artifact', 2:'Not Tracking', 3:'Tip Contamination'}
+    return class_names
+
+# Preprocess uploaded image
+def preprocess_image(img):
+    preprocess = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.3718, 0.1738, 0.0571], 
+            std=[0.2095, 0.2124, 0.1321]
+        ),
+    ])
+    img_tensor = preprocess(img).unsqueeze(0)  # Add batch dimension
+    return img_tensor
+
+# Predict using PyTorch model
+def predict_image_class(img, model, class_names):
+    img_tensor = preprocess_image(img)
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        _, preds = torch.max(outputs, dim=1)
+        print(_, preds)
+        probs = F.softmax(outputs, dim=1)
+        top_prob, top_idx = torch.topk(probs, 1)
+        print(top_prob, top_idx)
+        class_label = class_names[top_idx.item()]
+    return class_label
+
+img_path = '/work/mech-ai/angona3/Trial/image/Not_Tracking/Not_Tracking_21.jpg'
+# img_path = '/work/mech-ai/angona3/Trial/image/Tip_Contamination/Tip_Contamination_17.jpg'
+img = Image.open(img_path)
+model = load_vgg16()
+class_names = load_class_labels()
+class_label = predict_image_class(img, model, class_names)
+print(class_label)

graph_logic.py

@@ -0,0 +1,82 @@
+# --- START OF FILE graph_logic.py ---
+
+import os
+from typing import TypedDict, Annotated, List
+import operator
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+from langchain_anthropic import ChatAnthropic
+from langchain.schema import BaseMessage, AIMessage
+from langgraph.graph import StateGraph, END
+from langgraph.checkpoint.memory import MemorySaver # Optional for checkpointing
+
+# Load environment variables (optional here, but good practice if testing independently)
+# load_dotenv() # Can be commented out if only app_main.py loads it
+
+# --- LangGraph State Definition ---
+
+class GraphState(TypedDict):
+    """
+    Represents the state of our graph.
+
+    Attributes:
+        messages: The list of messages comprising the conversation.
+                  operator.add indicates messages should be appended.
+    """
+    messages: Annotated[List[BaseMessage], operator.add]
+
+# --- LLM Initialization ---
+
+def initialize_llm(provider: str, model_name: str, temperature: float, api_key: str):
+    """Initializes the appropriate LangChain Chat Model."""
+    if provider == "OpenAI":
+        if not api_key:
+            raise ValueError("OpenAI API key is missing. Please set OPENAI_API_KEY.")
+        return ChatOpenAI(api_key=api_key, model_name=model_name, temperature=temperature)
+    elif provider == "Anthropic":
+        if not api_key:
+            raise ValueError("Anthropic API key is missing. Please set ANTHROPIC_API_KEY.")
+        return ChatAnthropic(api_key=api_key, model_name=model_name, temperature=temperature)
+    else:
+        raise ValueError(f"Unsupported LLM provider: {provider}")
+
+# --- LangGraph Node and Graph Building ---
+
+def create_chat_graph(llm):
+    """
+    Builds and compiles the LangGraph conversational graph.
+
+    Args:
+        llm: An initialized LangChain Chat Model instance.
+
+    Returns:
+        A compiled LangGraph application.
+    """
+
+    # Define the function that calls the LLM - it closes over the 'llm' variable
+    def call_model(state: GraphState) -> dict:
+        """Invokes the provided LLM with the current conversation state."""
+        messages = state['messages']
+        response = llm.invoke(messages)
+        # Return the AIMessage list to be added to the state
+        return {"messages": [response]}
+
+    # Build the graph workflow
+    workflow = StateGraph(GraphState)
+
+    # Add the single node that runs the LLM
+    workflow.add_node("llm_node", call_model)
+
+    # Set the entry point and the only edge
+    workflow.set_entry_point("llm_node")
+    workflow.add_edge("llm_node", END) # Conversation ends after one LLM call per turn
+
+    # Compile the graph
+    # Optional: Add memory for checkpointing if needed
+    # memory = MemorySaver()
+    # graph = workflow.compile(checkpointer=memory)
+    graph = workflow.compile()
+
+    return graph
+
+# --- END OF FILE graph_logic.py ---

requirements.txt

@@ -0,0 +1,44 @@
+# Core LangChain packages
+langchain
+langchain-core
+
+# LangGraph
+langgraph
+
+# Pydantic
+pydantic
+
+# OpenAI integration (GPT-4, GPT-4o)
+langchain-openai
+openai
+
+# Anthropic integration (Claude 3)
+langchain-anthropic
+anthropic
+
+# Google Gemini integration (Gemini Pro / Flash)
+langchain-google-genai
+google-generativeai
+
+# Groq integration (Groq)
+langchain-groq
+groq
+
+# Data processing
+pandas
+PyPDF2
+
+# Optional: for improved dev type checking
+typing-extensions
+
+# Environment variables
+python-dotenv
+
+# for streamlit
+streamlit
+
+# for classification model
+torch
+torchvision
+numpy
+PIL