Spaces:
Sleeping
Sleeping
File size: 8,067 Bytes
70d5cce |
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 |
# from fastapi import FastAPI, File, UploadFile, HTTPException
# from fastapi.responses import JSONResponse
# import logging
# from ultralytics import YOLO
# import numpy as np
# import cv2
# from io import BytesIO
# from PIL import Image
# import base64
# import os
# # Setup logging
# logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
# logger = logging.getLogger(__name__)
# app = FastAPI(title="Car Parts & Damage Detection API")
# # Log model file presence
# model_files = ["car_part_detector_model.pt", "damage_general_model.pt"]
# for model_file in model_files:
# if os.path.exists(model_file):
# logger.info(f"Model file found: {model_file}")
# else:
# logger.error(f"Model file missing: {model_file}")
# # Load YOLO models
# try:
# logger.info("Loading car part model...")
# car_part_model = YOLO("car_part_detector_model.pt")
# logger.info("Car part model loaded successfully")
# logger.info("Loading damage model...")
# damage_model = YOLO("damage_general_model.pt")
# logger.info("Damage model loaded successfully")
# except Exception as e:
# logger.error(f"Failed to load models: {str(e)}")
# raise RuntimeError(f"Failed to load models: {str(e)}")
# def image_to_base64(img: np.ndarray) -> str:
# """Convert numpy image to base64 string."""
# try:
# _, buffer = cv2.imencode(".png", img)
# return base64.b64encode(buffer).decode("utf-8")
# except Exception as e:
# logger.error(f"Error encoding image to base64: {str(e)}")
# raise
# @app.post("/predict", summary="Run inference on an image for car parts and damage")
# async def predict(file: UploadFile = File(...)):
# """Upload an image and get car parts and damage detection results."""
# logger.info("Received image upload")
# try:
# contents = await file.read()
# image = Image.open(BytesIO(contents)).convert("RGB")
# img = np.array(image)
# logger.info(f"Image loaded: shape={img.shape}")
# blank_img = np.full((img.shape[0], img.shape[1], 3), 128, dtype=np.uint8)
# car_part_img = blank_img.copy()
# damage_img = blank_img.copy()
# car_part_text = "Car Parts: No detections"
# damage_text = "Damage: No detections"
# try:
# logger.info("Running car part detection...")
# car_part_results = car_part_model(img)[0]
# if car_part_results.boxes:
# car_part_img = car_part_results.plot()[..., ::-1]
# car_part_text = "Car Parts:\n" + "\n".join(
# f"- {car_part_results.names[int(cls)]} ({conf:.2f})"
# for conf, cls in zip(car_part_results.boxes.conf, car_part_results.boxes.cls)
# )
# logger.info("Car part detection completed")
# except Exception as e:
# car_part_text = f"Car Parts: Error: {str(e)}"
# logger.error(f"Car part detection error: {str(e)}")
# try:
# logger.info("Running damage detection...")
# damage_results = damage_model(img)[0]
# if damage_results.boxes:
# damage_img = damage_results.plot()[..., ::-1]
# damage_text = "Damage:\n" + "\n".join(
# f"- {damage_results.names[int(cls)]} ({conf:.2f})"
# for conf, cls in zip(damage_results.boxes.conf, damage_results.boxes.cls)
# )
# logger.info("Damage detection completed")
# except Exception as e:
# damage_text = f"Damage: Error: {str(e)}"
# logger.error(f"Damage detection error: {str(e)}")
# car_part_img_base64 = image_to_base64(car_part_img)
# damage_img_base64 = image_to_base64(damage_img)
# logger.info("Returning prediction results")
# return JSONResponse({
# "car_part_image": car_part_img_base64,
# "car_part_text": car_part_text,
# "damage_image": damage_img_base64,
# "damage_text": damage_text
# })
# except Exception as e:
# logger.error(f"Inference error: {str(e)}")
# raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")
# @app.get("/", summary="Health check")
# async def root():
# """Check if the API is running."""
# logger.info("Health check accessed")
# return {"message": "Car Parts & Damage Detection API is running"}
import gradio as gr
import numpy as np
import cv2
from PIL import Image
import base64
from io import BytesIO
from ultralytics import YOLO
import logging
import time
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Load damage detection model
try:
logger.info("Loading damage model...")
damage_model = YOLO("damage_general_model.pt")
logger.info("Damage model loaded successfully")
except Exception as e:
logger.error(f"Failed to load damage model: {str(e)}")
raise RuntimeError(f"Failed to load damage model: {str(e)}")
def image_to_base64(img: np.ndarray) -> str:
"""Convert numpy image to base64 string."""
try:
_, buffer = cv2.imencode(".png", img)
return base64.b64encode(buffer).decode("utf-8")
except Exception as e:
logger.error(f"Error encoding image to base64: {str(e)}")
raise
def process_images(*images):
"""Process up to 5 images for damage detection."""
if not any(images):
return "Please upload at least one image.", []
results = []
timing_info = []
for idx, img in enumerate(images):
if img is None:
continue
try:
start_image_time = time.time() # Start timer for individual image
logger.info(f"Processing image {idx + 1}")
# Convert Gradio image input (PIL) to numpy
img_np = np.array(img)
blank_img = np.full((img_np.shape[0], img_np.shape[1], 3), 128, dtype=np.uint8)
damage_text = f"Image {idx + 1} - Damage: No detections"
damage_img = blank_img.copy()
# Run damage detection
logger.info(f"Running damage detection for image {idx + 1}...")
damage_results = damage_model(img_np)[0]
if damage_results.boxes:
damage_img = damage_results.plot()[..., ::-1]
damage_text = f"Image {idx + 1} - Damage:\n" + "\n".join(
f"- {damage_results.names[int(cls)]} ({conf:.2f})"
for conf, cls in zip(damage_results.boxes.conf, damage_results.boxes.cls)
)
logger.info(f"Damage detection completed for image {idx + 1}")
# Convert result image to PIL for Gradio display
damage_pil = Image.fromarray(damage_img)
results.append((damage_pil, damage_text))
except Exception as e:
logger.error(f"Error processing image {idx + 1}: {str(e)}")
results.append((None, f"Image {idx + 1} - Error: {str(e)}"))
# Calculate total processing time
total_time = time.time() - start_image_time
timing_info.append(f"Total processing time: {total_time:.2f} seconds")
return "Damage detection completed.", results, "\n".join(timing_info)
# Define Gradio interface
iface = gr.Interface(
fn=process_images,
inputs=[
gr.Image(type="pil", label="Upload Image 1"),
gr.Image(type="pil", label="Upload Image 2"),
gr.Image(type="pil", label="Upload Image 3"),
gr.Image(type="pil", label="Upload Image 4"),
gr.Image(type="pil", label="Upload Image 5"),
],
outputs=[
gr.Textbox(label="Status"),
gr.Gallery(label="Detected Damage Images and Results", columns=2),
],
title="Car Damage Detection",
description="Upload up to 5 images to detect car damage. Results will display annotated images and detected damage details.",
)
# Launch the Gradio app
if __name__ == "__main__":
iface.launch() |