prolific / chatbot.py
Umar001's picture
Upload folder using huggingface_hub
5b11a78 verified
Raw
History Blame Contribute Delete
6.75 kB
import google.generativeai as genai_generative
import json
import time
# Configure API key
genai_generative.configure(api_key="AIzaSyA9jCSJ6HTwoDeEDb4mO7dmh15dhQ64Kqc")
# Defect classification mapping
issue_map = """
Defect Description: > Hairline multiple shallow cracks.Small multiple cracks that looks like spider web.
Defect Class: RCC - Shrinkage cracks
Defect Description: > Found on load-bearing members like column, beam, slab Have more depth then hairline crack Have width of more than 0.3 mm.
Defect Class: RCC - Crack
Defect Description: > Detaching layer of concrete along the reinforcement / TMT. Exposed rusted reinforcement / TMT. Hollowness on concrete surface.
Defect Class: RCC - Spalling of concrete
Defect Description: > Hairline multiple shallow cracks. Small multiple cracks that looks like spider web.
Defect Class: Plaster - Shrinkage Cracks
Defect Description: > Found on RCC & brick / block wall junction Generally a straight crack (Vertical or horizontal)
Defect Class: Plaster - Separation Cracks
Defect Description: > Generally found on the middle of the wall. Either stepped or diagonal is nature.
Defect Class: Plaster - Diagonal crack with less than 3 mm width
Defect Description: > Generally found on the middle of the wall. Either stepped or diagonal is nature.
Defect Class: Plaster - Diagonal crack with more than 3 mm width
Defect Description: >Looks like uneven & patchy at places. Uneven gap noted between plaster & straight edge.
Defect Class: Plaster - Uneven Surface
Defect Description: > Loose sand comes out of plaster when rubbed with solid object.
Defect Class: Plaster - Loose sand
Defect Description: > Appear like loose white powder over plaster surface. Generally appear on moist / wet surface. Common with bed bricks walls.
Defect Class: Plaster - Efflorescence
Defect Description: > Hollow sound comes out of plaster when tapped with metal object / hollow stick. Delaminating of plaster layer from brick or RCC surface.
Defect Class: Plaster - Deboning (Hollowness)
"""
# Prompt builders
def _build_prompt(image_description: str = "") -> str:
if image_description:
return (
"Based on the image and image description, match the defect description, "
"and classify the defect into the correct defect class.\n\n"
f"Image Description: {image_description}\n\n"
"Defect Description and class: \n"
f"{issue_map}\n\n"
"format the response in the following format:\n"
"Defect Class: <defect_class>\n"
"Reasoning: <reasoning>"
)
return (
" Based on the image and the defect description, \n"
"please classify the defect into the correct defect class.\n"
"Defect Description and class: \n\n"
f"{issue_map}\n\n"
"format the response in the following format:\n"
"Defect Class: <defect_class>\n"
"Reasoning: <reasoning>"
)
# Initialize model
model = genai_generative.GenerativeModel("gemini-2.0-flash", generation_config={"temperature": 0.5})
def upload_file(image_path, image_description: str = ""):
"""Upload image to Gemini and get classification response"""
myfile = genai_generative.upload_file(image_path)
while myfile.state.name == "PROCESSING":
time.sleep(0.5)
myfile = genai_generative.get_file(myfile.name)
try:
prompt = _build_prompt(image_description)
image_response = model.generate_content(contents=[prompt, myfile])
print(image_response.text)
finally:
# Clean up the uploaded file to prevent gRPC timeout warnings
try:
genai_generative.delete_file(myfile.name)
except Exception as e:
print(f"Warning: Could not delete uploaded file: {e}")
return image_response.text
# Load GKA.json knowledge base
with open('GKA.json', 'r') as file:
gka_data = json.load(file)
gka_data_str = json.dumps(gka_data, indent=2)
def chatbot(image_path, image_description: str = ""):
"""Simple chatbot with max 2 interactions: image analysis + 1 follow-up question"""
print("=== Building Defect Analysis Chatbot ===")
print("Step 1: Analyzing image for defects...")
# Step 1: Analyze image
image_response = upload_file(image_path, image_description)
print("\n=== Image Analysis Complete ===")
print("Defect identified! You can ask 1 follow-up question.")
# Step 2: One follow-up question
query = input("\nEnter your question about this defect: ")
if query.strip():
print("\n=== Answering your question ===")
# Extract defect class and reasoning from image response
try:
defect_class = image_response.split("Defect Class: ")[1].split("\n")[0] if "Defect Class: " in image_response else "Unknown"
reasoning = image_response.split("Reasoning: ")[1] if "Reasoning: " in image_response else "No reasoning provided"
except:
defect_class = "Unknown"
reasoning = "No reasoning provided"
# Generate answer based on knowledge base
prompt = f"""Based on the defect analysis, answer this question using the knowledge base.
Defect Class: {defect_class}
Reasoning: {reasoning}
Question: {query}
Knowledge Base: {gka_data_str}
Provide a helpful answer based on the knowledge base information."""
answer = model.generate_content(contents=[prompt])
print(f"\nAnswer: {answer.text}")
return {
"defect_class": defect_class,
"reasoning": reasoning,
"question": query,
"answer": answer.text
}
else:
print("No question asked. Chatbot session complete.")
return {
"defect_class": image_response.split("Defect Class: ")[1].split("\n")[0] if "Defect Class: " in image_response else "Unknown",
"reasoning": image_response.split("Reasoning: ")[1] if "Reasoning: " in image_response else "No reasoning provided",
"question": None,
"answer": None
}
if __name__ == "__main__":
# Simple chatbot usage
image_path = "1000039877.jpeg"
result = chatbot(image_path)
print("\n=== Chatbot Session Summary ===")
print(f"Defect Class: {result['defect_class']}")
print(f"Reasoning: {result['reasoning']}")
if result['question']:
print(f"Question: {result['question']}")
print(f"Answer: {result['answer']}")
else:
print("No follow-up question was asked.")