Spaces:
Runtime error
Runtime error
File size: 3,346 Bytes
7b813cc 2204e0b 7b813cc d86c7bd |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
from flask import Flask, request, jsonify, render_template
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.data import MetadataCatalog
from detectron2.utils.visualizer import Visualizer, ColorMode
import numpy as np
from PIL import Image
import io
import os
import requests
import gdown
# Initialize Flask app
app = Flask(__name__)
cfg = None
# Google Drive file URL
GDRIVE_MODEL_URL = "https://drive.google.com/uc?id=18aEDo-kWOBhg8mAhnbpFkuM6bmmrBH4E" # Replace 'your-file-id' with the actual file ID from Google Drive
LOCAL_MODEL_PATH = "model_final.pth"
def download_file_from_google_drive(id, destination):
gdown.download(GDRIVE_MODEL_URL, LOCAL_MODEL_PATH, quiet=False)
file_id = '18aEDo-kWOBhg8mAhnbpFkuM6bmmrBH4E'
destination = 'model_final.pth'
download_file_from_google_drive(file_id, destination)
# Download model from Google Drive if not already present locally
def download_model():
if not os.path.exists(LOCAL_MODEL_PATH):
response = requests.get(GDRIVE_MODEL_URL, stream=True)
if response.status_code == 200:
with open(LOCAL_MODEL_PATH, 'wb') as f:
f.write(response.content)
else:
raise Exception(f"Failed to download model from Google Drive: {response.status_code}")
# Configuration and model setup
def setup_model(model_path):
global cfg
cfg = get_cfg()
cfg.merge_from_file("config.yaml") # Update with the config file path
cfg.MODEL.WEIGHTS = model_path
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.DEVICE = "cpu" # Use "cuda" for GPU
return DefaultPredictor(cfg)
# Ensure model is available
predictor = setup_model(LOCAL_MODEL_PATH)
# Define expected parts and costs
expected_parts = ['headlamp', 'rear_bumper', 'door', 'hood', 'front_bumper']
cost_dict = {
'headlamp': 300,
'rear_bumper': 250,
'door': 200,
'hood': 220,
'front_bumper': 250,
'other': 150
}
@app.route('/')
def home():
return render_template('index.html')
@app.route('/upload', methods=['POST'])
def upload():
if 'file' not in request.files:
return jsonify({"error": "No file uploaded"}), 400
file = request.files['file']
if file.filename == '':
return jsonify({"error": "No file selected"}), 400
# Load image
image = Image.open(file).convert("RGB")
image_np = np.array(image)
# Run model prediction
outputs = predictor(image_np)
instances = outputs["instances"].to("cpu")
class_names = MetadataCatalog.get(cfg.DATASETS.TEST[0]).thing_classes
# Initialize total cost
total_cost = 0
damage_details = []
for j in range(len(instances)):
class_id = instances.pred_classes[j].item()
damaged_part = class_names[class_id] if class_id < len(class_names) else 'unknown'
if damaged_part not in expected_parts:
damaged_part = 'other'
repair_cost = cost_dict.get(damaged_part, cost_dict['other'])
total_cost += repair_cost
damage_details.append({
'part': damaged_part,
'cost_usd': repair_cost
})
response = {
"damages": damage_details,
"total_cost": total_cost
}
return jsonify(response)
if __name__ == '__main__':
app.run(host='0.0.0.0', port=7860)
|