Spaces:
Sleeping
Sleeping
| # 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() |