File size: 2,622 Bytes
7136f34 83437a6 7136f34 7bbf949 7136f34 7bbf949 7136f34 d053647 7136f34 d053647 7136f34 d053647 7136f34 83437a6 7136f34 7bbf949 7136f34 d053647 7136f34 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
from flask import Flask, request, jsonify
import torch
from transformers import AutoProcessor, AutoModelForImageClassification
from sentence_transformers import SentenceTransformer
import sqlite3
app = Flask(__name__)
# Load the defect detection model (open-source, hugging face model)
DETECTION_MODEL_NAME = "microsoft/beit-base-patch16-224-pt22k-ft22k"
processor = AutoProcessor.from_pretrained(DETECTION_MODEL_NAME)
detection_model = AutoModelForImageClassification.from_pretrained(DETECTION_MODEL_NAME)
defects_to_remedies = {
"crack": "Fill cracks with epoxy. Structural cracks might need professional inspection.",
"spalling": "Clean affected area and apply anti-corrosion primer before repairing.",
"leakage": "Fix water source, seal with water-proofing compounds.",
"mold": "Clean the mold, improve ventilation, and apply mold-resistant paint."
}
# Initialize a Sentence Transformer for text embeddings
EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
# SQLite database setup
db_conn = sqlite3.connect('defects.db', check_same_thread=False)
c = db_conn.cursor()
c.execute('''CREATE TABLE IF NOT EXISTS defects (id INTEGER PRIMARY KEY, defect TEXT, remedy TEXT, embedding BLOB)''')
db_conn.commit()
# Populate defect remedies table if empty
def seed_database():
for defect, remedy in defects_to_remedies.items():
c.execute("SELECT * FROM defects WHERE defect=?", (defect,))
if not c.fetchone():
embedding = embedding_model.encode(remedy).tolist()
c.execute("INSERT INTO defects (defect, remedy, embedding) VALUES (?, ?, ?)", (defect, remedy, str(embedding)))
db_conn.commit()
seed_database()
@app.route('/detect', methods=['POST'])
def detect_defect():
if 'image' not in request.files:
return jsonify({"error": "No image uploaded."}), 400
image = request.files['image'].read()
# Preprocess and predict the defect
inputs = processor(images=image, return_tensors="pt")
outputs = detection_model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=1)
predicted_class = torch.argmax(probs, dim=1)
class_name = detection_model.config.id2label[predicted_class.item()]
# Query remedy
c.execute("SELECT remedy FROM defects WHERE defect=?", (class_name,))
row = c.fetchone()
if row:
remedy = row[0]
else:
remedy = "No specific remedy available for this defect."
return jsonify({"detected_defect": class_name, "remedy": remedy})
if __name__ == '__main__':
app.run(debug=True) |