Application update

app.py

@@ -7,15 +7,10 @@ Original file is located at
     https://colab.research.google.com/drive/1HCITtw0z2BJO0z9GmsspLM1NzrLdft27
 """
 
-#!pip install gradio --quiet
-#!pip install -Uq transformers datasets timm accelerate evaluate
-
-import subprocess
-# subprocess.run('pip3 install datasets timm cv2 huggingface_hub torch pillow matplotlib' ,shell=True)
-
 import gradio as gr
-from huggingface_hub import hf_hub_download
 from safetensors.torch import load_model
+from timm import create_model
+from huggingface_hub import hf_hub_download
 from datasets import load_dataset
 import torch
 import torchvision.transforms as T
@@ -23,8 +18,17 @@ import cv2
 import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image
-from timm import create_model
+import os
 
+from langchain_community.document_loaders import TextLoader
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langchain_fireworks import ChatFireworks
+from langchain_core.prompts import ChatPromptTemplate
+from transformers import AutoModelForImageClassification, AutoImageProcessor
 
 
 safe_tensors = "model.safetensors" #hf_hub_download(repo_id="subhuatharva/swim-224-base-satellite-image-classification", filename="model.safetensors")
@@ -52,8 +56,56 @@ def one_hot_decoding(labels):
     true_labels.append(id2label[i])
   return true_labels
 
+def ragChain():
+    """
+    function: creates a rag chain
+    output: rag chain
+    """
+    loader = TextLoader("document.txt")
+    docs = loader.load()
+
+    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+    docs = text_splitter.split_documents(docs)
+
+    vectorstore = FAISS.load_local("faiss_index", embeddings = HuggingFaceEmbeddings(), allow_dangerous_deserialization = True)
+    retriever = vectorstore.as_retriever(search_type = "similarity", search_kwargs = {"k": 5})
+
+    api_key = os.getenv("FIREWORKS_API_KEY")
+    llm = ChatFireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct", api_key = api_key)
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                """You are a knowledgeable landscape deforestation analyst.
+                 """
+            ),
+            (
+                "human",
+                """First mention the detected labels only with short description.
+                Provide not more than 4 precautionary measures which are related to the detected labels that can be taken to control deforestation.
+                Don't include conversational messages.
+                """,
+            ),
+
+            ("human", "{context}, {question}"),
+        ]
+    )
+
+    rag_chain = (
+    {
+    "context": retriever,
+    "question": RunnablePassthrough()
+    }
+    | prompt
+    | llm
+    | StrOutputParser()
+    )
+
+    return rag_chain
+
 def model_output(image):
-  image = cv2.imread(image)
+
   PIL_image = Image.fromarray(image.astype('uint8'), 'RGB')
 
   img_size = (224,224)
@@ -72,8 +124,27 @@ def model_output(image):
   pred_labels = one_hot_decoding(predictions)
   output_text = " ".join(pred_labels)
 
-  return output_text
-
-app = gr.Interface(fn=model_output, inputs="image", outputs="text")
-app.launch(debug=True)
+  query = f"Detected labels in the provided satellite image are {output_text}. Give information on the labels."
+
+  return query
+
+def generate_response(rag_chain, query):
+    """
+    input: rag chain, query
+    function: generates response using llm and knowledge base
+    output: generated response by the llm
+    """
+    return rag_chain.invoke(f"{query}")
+    
+def main(image):
+    query = model_output(image)
+    chain = ragChain()
+    output = generate_response(chain, query)
+    return output
+title = "Satellite Image Landscape Analysis for Deforestation"
+description = "This bot will take any satellite image and analyze the factors which lead to deforestation by identify the landscape based on forest areas, roads, habitation, water etc."
+app = gr.Interface(fn=main, inputs="image", outputs="text", title=title,
+                   description=description,
+                   examples=[["sampleimages/train_142.jpg"], ["sampleimages/train_32.jpg"],["sampleimages/train_59.jpg"], ["sampleimages/train_67.jpg"],["sampleimages/train_75.jpg"],["sampleimages/train_92.jpg"],["sampleimages/random_satellite.jpg"]])
+app.launch(share = True)
 

requirements.txt

@@ -1,6 +1,6 @@
 transformers
 datasets 
-timm 
+Time 
 langchain-fireworks 
 langchain_core 
 langchain_community 
@@ -10,4 +10,4 @@ safetensors
 torch
 torchvision
 opencv-python
-pillow
+pillow