Update app.py

app.py

@@ -14,7 +14,11 @@ from langchain.agents.agent_types import AgentType
 from langchain.agents import Tool, initialize_agent
 from langchain_community.callbacks.streamlit import StreamlitCallbackHandler
 from groq import Groq
+import open_clip
+from open_clip import create_model_from_pretrained, get_tokenizer # works on open-clip-torch>=2.23.0, timm>=0.9.8
 
+model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224')
+tokenizer = open_clip.get_tokenizer('hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224')
 load_dotenv()
 groq_api_key = os.getenv("GROQ_API_KEY")
 
@@ -58,10 +62,33 @@ reasoning_tool = Tool(
     func=chain.run,
     description="A tool for answering logic-based and reasoning questions."
 )
-
+def classify_image(image_path: str) -> str:
+    """Classifies a medical image using BiomedCLIP."""
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    model.to(device).eval()
+    
+    # Open and preprocess image
+    image = preprocess(Image.open(image_path)).unsqueeze(0).to(device)
+    labels = ["MRI scan", "X-ray", "histopathology", "CT scan", "ultrasound", "medical chart"]
+    texts = tokenizer([f"this is a photo of {l}" for l in labels], context_length=256).to(device)
+    
+    with torch.no_grad():
+        image_features, text_features, logit_scale = model(image, texts)
+        logits = (logit_scale * image_features @ text_features.t()).detach().softmax(dim=-1)
+        sorted_indices = torch.argsort(logits, dim=-1, descending=True)
+    
+    top_class = labels[sorted_indices[0][0].item()]
+    return f"The image is classified as {top_class}."
+
+# Wrap BiomedCLIP as a LangChain tool
+biomed_clip_tool = Tool(
+    name="BiomedCLIP Image Classifier",
+    func=classify_image,
+    description="Classifies medical images into categories like MRI, X-ray, histopathology, etc."
+)
 # Initialize the agents for text questions
 assistant_agent_text = initialize_agent(
-    tools=[wikipedia_tool, calculator, reasoning_tool],
+    tools=[wikipedia_tool, calculator, reasoning_tool, biomed_clip_tool],
     llm=llm_text,
     agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
     verbose=False,
@@ -168,4 +195,4 @@ elif st.session_state["section"] == "image":
                 except ValueError as e:
                     st.error(f"An error occurred: {e}")
         else:
-            st.warning("Please enter a question and upload an image to get an answer.")
+            st.warning("Please enter a question and upload an image to get an answer.")