Filupa/Object_Detection

1import os2import io3import logging4from fastapi import FastAPI, File, UploadFile, HTTPException5from fastapi.responses import FileResponse6from fastapi.middleware.cors import CORSMiddleware7from PIL import Image, ExifTags8from ultralytics import YOLO9import numpy as np10import cv211 12# Direktori penyimpanan sementara di /tmp/13IMAGE_PATH = "/tmp/processed_image.jpg"14 15server = FastAPI()16 17# Enable CORS18server.add_middleware(19    CORSMiddleware,20    allow_origins=["*"],21    allow_credentials=True,22    allow_methods=["*"],23    allow_headers=["*"],24)25 26CLASS_NAMES = {27    0: "Ice Cream",28    1: "Lollipop",29    2: "Chocolate",30    3: "Train",31    4: "Minibus",32    5: "Plane",33    6: "Bee",34    7: "Sheep",35    8: "Cat",36    9: "Strawberry",37    10: "Banana",38    11: "Grape"39}40 41# Load YOLO model42try:43    model = YOLO("best.pt")44    print("✅ YOLO model loaded successfully")45except Exception as e:46    print(f"❌ Model loading failed: {e}")47    raise RuntimeError("Model failed to load")48 49def fix_image_rotation(image_bytes):50    image = Image.open(io.BytesIO(image_bytes))51 52    try:53        for orientation in ExifTags.TAGS.keys():54            if ExifTags.TAGS[orientation] == 'Orientation':55                break56 57        exif = image._getexif()58        if exif is not None:59            orientation = exif.get(orientation, 1)60            if orientation == 3:61                image = image.rotate(180, expand=True)62            elif orientation == 6:63                image = image.rotate(270, expand=True)64            elif orientation == 8:65                image = image.rotate(90, expand=True)66 67    except (AttributeError, KeyError, IndexError):68        pass69 70    return image71 72@server.get("/")73async def root():74    return {"message": "Server is running"}75 76@server.post("/predict/")77async def predict(file: UploadFile = File(...)):78    try:79        # Baca gambar80        contents = await file.read()81        image = fix_image_rotation(contents)82        image = np.array(image.convert("RGB"))83        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)84 85        # Letterbox Resize Function86        def letterbox(im, new_shape=(640, 640), color=(114, 114, 114)):87            shape = im.shape[:2]88            ratio = min(new_shape[0] / shape[0], new_shape[1] / shape[1])89            new_unpad = int(round(shape[1] * ratio)), int(round(shape[0] * ratio))90            im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)91            dh, dw = new_shape[0] - new_unpad[1], new_shape[1] - new_unpad[0]92            dh, dw = dh // 2, dw // 293            im = cv2.copyMakeBorder(im, dh, dh, dw, dw, cv2.BORDER_CONSTANT, value=color)94            return im95 96        # Resize gambar97        image = letterbox(image, (640, 640))98 99        # Jalankan deteksi YOLO100        results = model(image)101 102        # Ambil prediksi terbaik103        best_prediction = None104        105        for result in results:106            for box in result.boxes:107                x1, y1, x2, y2 = map(int, box.xyxy[0])108                confidence = float(box.conf[0])  # Ambil confidence score109                class_id = int(box.cls[0])  # Ambil class ID110        111                # Ambil nama kelas dari dictionary112                class_name = CLASS_NAMES.get(class_id, f"Unknown ({class_id})")113        114                # Pilih prediksi terbaik berdasarkan confidence tertinggi115                if best_prediction is None or confidence > best_prediction["confidence"]:116                    best_prediction = {"class": class_name, "confidence": confidence}117        118                # Gambar bounding box119                cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)120        121        # Simpan gambar hasil deteksi ke direktori /tmp/122        cv2.imwrite(IMAGE_PATH, image)123        124        # Jika tidak ada objek yang terdeteksi, kirimkan respons kosong125        if best_prediction is None:126            return {"message": "No object detected"}127 128        print(f"Best Prediction: {best_prediction['class']} | Confidence: {best_prediction['confidence']:.4f}")129        130        # Kirim respons hasil deteksi131        response_data = {132            "prediction": best_prediction["class"],  # Mengirim nama kelas, bukan ID133            "confidence": round(best_prediction["confidence"], 2),134        }135        136        return response_data137 138        139 140    except Exception as e:141        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")142 143@server.get("/image/")144async def get_image():145    if os.path.exists(IMAGE_PATH):146        return FileResponse(IMAGE_PATH, media_type="image/jpeg")147    raise HTTPException(status_code=404, detail="Image not found!")