Swagat23154's picture
Upload app.py
c1cb7d8 verified
from ultralytics import YOLO
import cv2
import torch
import gradio as gr
from PIL import Image
import numpy as np
import openai
import base64
import io
import os
# Load YOLOv8 model (Ensure best.pt is in the same directory)
model_path = "best.pt"
if not os.path.exists(model_path):
raise FileNotFoundError(f"Model file '{model_path}' not found. Upload it to the Space.")
model = YOLO(model_path) # Load model
model.to("cpu") # Ensure it runs on CPU
# OpenAI API Key (Set as environment variable in Hugging Face Space)
openai_api_key = os.getenv('openai_api_key',"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")
print(f"DEBUG: OpenAI API Key - {openai_api_key[:5]}**********") # Prints only first 5 characters for security
openai.api_key = openai_api_key
def encode_image(image):
"""Convert image to base64 for OpenAI API"""
buffered = io.BytesIO()
image.save(buffered, format="JPEG")
return base64.b64encode(buffered.getvalue()).decode()
def get_car_details(image):
"""Send the image to OpenAI to get car make, model, and year."""
try:
base64_image = encode_image(image) # Convert image to Base64
openai_api_key = os.getenv('openai_api_key',"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")
if not openai_api_key:
raise ValueError("❌ OpenAI API key is missing! Set it in the environment.")
client = openai.OpenAI(api_key=openai_api_key) # ✅ Pass API key explicitly
response = client.chat.completions.create(
model="gpt-4o-mini", # 🔹 Use GPT-4o Mini
messages=[
{"role": "system", "content": "You are an AI that identifies car make, model, and year from images."},
{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
{"type": "text", "text": "Identify the make, model, and year of this car."}
]}
],
max_tokens=100
)
return response.choices[0].message.content.strip()
except Exception as e:
return f"Error getting car details: {str(e)}"
def predict(input_img):
# Convert PIL image to OpenCV format
image = np.array(input_img)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# Run inference
results = model(image)
# Confidence threshold
threshold = 0.4 # Show only if confidence > 0.55
# Draw bounding boxes
for result in results:
for box in result.boxes:
conf = float(box.conf[0]) # Confidence score
if conf > threshold: # ✅ Only consider if confidence > 0.55
x1, y1, x2, y2 = map(int, box.xyxy[0])
label = f"Damage: {conf:.2f}"
# Draw red bounding box
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 3)
cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,
0.5, (0, 0, 255), 2)
# Convert back to PIL format
output_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
# Get car details from OpenAI
car_details = get_car_details(input_img)
return output_img, car_details
# Gradio interface
gradio_app = gr.Interface(
fn=predict,
inputs=gr.Image(label="Upload a car image", sources=['upload', 'webcam'], type="pil"),
outputs=[
gr.Image(label="Detected Damage"),
gr.Textbox(label="Car Make, Model & Year")
],
title="Car Damage & Identification",
description="Upload an image of a car to detect damage and identify the make, model, and year."
)
if __name__ == "__main__":
gradio_app.launch(share=True)